Struct PsaPackQkvOp 

pub struct PsaPackQkvOp<D: Float + Send + Sync + 'static> { /* private fields */ }

CustomOp for PsaPackQkvRuntimeOp.

Graph node: [B, qkv_h, H, W] → [4, BH, N, KEY_DIM]

Implementations

impl<D: Float + Send + Sync + 'static> PsaPackQkvOp<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 PsaPackQkvOp<D>

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

Output shape: [B, 4, num_heads, N, KEY_DIM]

B is kept as the leading (potentially-None) dimension so that resolve_shape(…, batch_size) sets it to batch_size directly. Placing B first avoids the prior bug where BH = B * num_heads was stored as None and resolved to batch_size instead of batch_size * num_heads.

fn name(&self) -> &str

Identifier used in error messages and debug output.

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.