Trait CustomOp 

pub trait CustomOp:
    Any
    + Send
    + Sync {
    // Required methods
    fn name(&self) -> &str;
    fn infer_output_shape(&self, input_shapes: &[&Shape]) -> Shape;
    fn as_any(&self) -> &dyn Any;

    // Provided methods
    fn lower(&self) -> Option<(String, String, String, Arc<dyn RuntimeOp>)> { ... }
    fn lower_backward_source(&self) -> String { ... }
}

Trait implemented by user-defined ops.

Required Methods

fn name(&self) -> &str

Identifier used in error messages and debug output.

fn infer_output_shape(&self, input_shapes: &[&Shape]) -> Shape

Compute the output shape given the shapes of all input tensors in order.

fn as_any(&self) -> &dyn Any

Expose self as &dyn Any so the custom lowering can downcast to the concrete op type. Implement as fn as_any(&self) -> &dyn Any { self }.

Provided Methods

fn lower(&self) -> Option<(String, String, String, Arc<dyn RuntimeOp>)>

Return kernel lowering info so TritonLowering can compile this op without a project-specific middleware. Return None to keep the existing middleware / error behaviour.

Tuple layout: (name, kernel_source, entry_point_name, runtime_op). entry_point_name is the PTX symbol name, conventionally "{name}_entry_point".

fn lower_backward_source(&self) -> String

Return the backward kernel source for this op (used in training mode). Return an empty string if this op has no backward pass.

Dyn Compatibility

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors