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.

Implementors