Struct PsaExtractVOp 

pub struct PsaExtractVOp<D: Float + Send + Sync + 'static>(/* private fields */);

CustomOp for PsaExtractVRuntimeOp.

Graph node: [B, qkv_h, H, W] → [B, c, H, W] (c = qkv_h / 2)

Implementations

impl<D: Float + Send + Sync + 'static> PsaExtractVOp<D>

pub fn new(key_dim: i32, num_heads: usize) -> Self

Creates the graph op for the given key dimension and head count.

Trait Implementations

impl<D: Float + Send + Sync + 'static> CustomOp for PsaExtractVOp<D>

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 }.

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.

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.