oa::ml::matrix

Module published by the oa::ml module.

MdRust public API

Public Types

Public Functions

pub fn adaptive_avg_pool_2d( input: &Matrix, output_height: usize, output_width: usize, ) -> Result<Matrix>

pub fn avg_pool_2d( input: &Matrix, kernel_size: usize, stride: usize, padding: usize, ) -> Result<Matrix>

pub fn batch_norm_2d( input: &Matrix, weight: &Matrix, bias: &Matrix, epsilon: f32, ) -> Result<BatchNorm2dResult>

pub fn batch_norm_2d_with_stats( input: &Matrix, mean: &Matrix, variance: &Matrix, weight: &Matrix, bias: &Matrix, epsilon: f32, ) -> Result<Matrix>

pub fn bmm(left: &Matrix, right: &Matrix) -> Result<Matrix>

pub fn bmm_nt(left: &Matrix, right: &Matrix) -> Result<Matrix>

pub fn bmm_tn(left: &Matrix, right: &Matrix) -> Result<Matrix>

pub fn channel_norm( input: &Matrix, weight: &Matrix, bias: &Matrix, epsilon: f32, ) -> Result<Matrix>

pub fn channel_norm_backward( input: &Matrix, weight: &Matrix, output_gradient: &Matrix, epsilon: f32, ) -> Result<ChannelNormBackward>

pub fn channel_norm_relu( input: &Matrix, weight: &Matrix, bias: &Matrix, epsilon: f32, ) -> Result<Matrix>

pub fn channel_norm_relu_backward( input: &Matrix, weight: &Matrix, forward_output: &Matrix, output_gradient: &Matrix, epsilon: f32, ) -> Result<ChannelNormBackward>

pub fn conv_1d( input: &Matrix, weight: &Matrix, bias: &Matrix, stride: usize, padding: usize, dilation: usize, ) -> Result<Matrix>

pub fn conv_2d( input: &Matrix, weight: &Matrix, bias: &Matrix, stride: usize, padding: usize, groups: usize, ) -> Result<Matrix>

pub fn conv_transpose_1d( input: &Matrix, weight: &Matrix, stride: usize, padding: usize, dilation: usize, ) -> Result<Matrix>

pub fn conv_transpose_2d( input: &Matrix, weight: &Matrix, bias: &Matrix, stride: usize, padding: usize, ) -> Result<Matrix>

pub fn detach(input: &Matrix) -> Result<Matrix>

pub fn elu(input: &Matrix, alpha: f32) -> Result<Matrix>

pub fn flash_attention_causal( query: &Matrix, key: &Matrix, value: &Matrix, scale: f32, ) -> Result<Matrix>

pub fn gelu(input: &Matrix) -> Result<Matrix>

pub fn grouped_gemm_m( input: &Matrix, weight: &Matrix, offsets: &Matrix, ) -> Result<Matrix>

pub fn grouped_linear_m( input: &Matrix, weight: &Matrix, bias: &Matrix, offsets: &Matrix, ) -> Result<Matrix>

pub fn gru_cell( gates_i: &Matrix, hidden: &Matrix, weight_hh: &Matrix, bias_hh: Option<&Matrix>, ) -> Result<Matrix>

pub fn gru_scan( gates_i: &Matrix, weight_hh: &Matrix, bias_hh: Option<&Matrix>, ) -> Result<Matrix>

pub fn leaky_relu(input: &Matrix, alpha: f32) -> Result<Matrix>

pub fn mamba3_mimo( c: &Matrix, b: &Matrix, x: &Matrix, z: &Matrix, adt: &Matrix, dt: &Matrix, trap: &Matrix, angle: &Matrix, c_bias: &Matrix, b_bias: &Matrix, d: &Matrix, mimo_x: &Matrix, mimo_z: &Matrix, mimo_o: &Matrix, norm_weight: &Matrix, config: SsmConfig, ) -> Result<Matrix>

pub fn mamba3_mimo_backward( output_gradient: &Matrix, c: &Matrix, b: &Matrix, x: &Matrix, z: &Matrix, adt: &Matrix, dt: &Matrix, trap: &Matrix, angle: &Matrix, c_bias: &Matrix, b_bias: &Matrix, d: &Matrix, mimo_x: &Matrix, mimo_z: &Matrix, mimo_o: &Matrix, norm_weight: &Matrix, config: SsmConfig, ) -> Result<Mamba3MimoBackward>

pub fn mamba3_mimo_step( c: &Matrix, b: &Matrix, x: &Matrix, z: &Matrix, adt: &Matrix, dt: &Matrix, trap: &Matrix, angle: &Matrix, c_bias: &Matrix, b_bias: &Matrix, d: &Matrix, mimo_x: &Matrix, mimo_z: &Matrix, mimo_o: &Matrix, norm_weight: &Matrix, ssm_state: &mut Matrix, angle_state: &mut Matrix, k_state: &mut Matrix, v_state: &mut Matrix, config: SsmConfig, ) -> Result<Matrix>

pub fn mamba3_preprocess( projected: &Matrix, dt_bias: &Matrix, config: Mamba3PreprocessConfig, ) -> Result<Mamba3PreprocessResult>

pub fn mamba3_preprocess_backward( projected: &Matrix, dt_bias: &Matrix, output_gradients: &Mamba3PreprocessResult, config: Mamba3PreprocessConfig, ) -> Result<Mamba3PreprocessBackward>

pub fn mamba3_siso( c: &Matrix, b: &Matrix, x: &Matrix, z: &Matrix, adt: &Matrix, dt: &Matrix, trap: &Matrix, angle: &Matrix, c_bias: &Matrix, b_bias: &Matrix, d: &Matrix, config: SsmConfig, ) -> Result<Matrix>

pub fn mamba3_siso_backward( output_gradient: &Matrix, c: &Matrix, b: &Matrix, x: &Matrix, z: &Matrix, adt: &Matrix, dt: &Matrix, trap: &Matrix, angle: &Matrix, c_bias: &Matrix, b_bias: &Matrix, d: &Matrix, config: SsmConfig, ) -> Result<SsmBackward>

pub fn mamba3_siso_step( c: &Matrix, b: &Matrix, x: &Matrix, z: &Matrix, adt: &Matrix, dt: &Matrix, trap: &Matrix, angle: &Matrix, c_bias: &Matrix, b_bias: &Matrix, d: &Matrix, ssm_state: &mut Matrix, angle_state: &mut Matrix, k_state: &mut Matrix, v_state: &mut Matrix, config: SsmConfig, ) -> Result<Matrix>

pub fn max_pool_2d( input: &Matrix, kernel_size: usize, stride: usize, padding: usize, ) -> Result<MaxPool2dResult>

pub fn merge_heads( input: &Matrix, batch: usize, sequence_length: usize, num_heads: usize, ) -> Result<Matrix>

pub fn mish(input: &Matrix) -> Result<Matrix>

pub fn moe_combine( packed: &Matrix, route_gate: &Matrix, inverse: &Matrix, packed_slot: &Matrix, ) -> Result<Matrix>

pub fn moe_gather( input: &Matrix, packed_token: &Matrix, inverse: &Matrix, ) -> Result<Matrix>

pub fn moe_route_weights( probabilities: &Matrix, expert_indices: &Matrix, ) -> Result<Matrix>

pub fn relu(input: &Matrix) -> Result<Matrix>

pub fn rms_norm(input: &Matrix, weight: &Matrix, epsilon: f32) -> Result<Matrix>

pub fn rms_norm_gated( input: &Matrix, weight: &Matrix, bias: Option<&Matrix>, gate: &Matrix, epsilon: f32, ) -> Result<Matrix>

pub fn rms_norm_gated_backward( input: &Matrix, weight: &Matrix, bias: Option<&Matrix>, gate: &Matrix, output_gradient: &Matrix, epsilon: f32, ) -> Result<RmsNormGatedBackward>

pub fn rnn_cell( gates_i: &Matrix, hidden: &Matrix, weight_hh: &Matrix, bias_hh: Option<&Matrix>, ) -> Result<Matrix>

pub fn rnn_scan( gates_i: &Matrix, weight_hh: &Matrix, bias_hh: Option<&Matrix>, ) -> Result<Matrix>

pub fn rope( input: &Matrix, num_heads: usize, head_dim: usize, theta_base: f32, position_offset: u32, ) -> Result<Matrix>

pub fn scaled_dot_product_attention( query: &Matrix, key: &Matrix, value: &Matrix, additive_mask: Option<&Matrix>, scale: f32, causal: bool, ) -> Result<Matrix>

pub fn scaled_dot_product_attention_causal( query: &Matrix, key: &Matrix, value: &Matrix, sequence_length: usize, num_heads: usize, ) -> Result<Matrix>

pub fn sigmoid(input: &Matrix) -> Result<Matrix>

pub fn silu(input: &Matrix) -> Result<Matrix>

pub fn silu_mul(input: &Matrix, intermediate_size: usize) -> Result<Matrix>

pub fn softmax_scaled_masked( scores: &Matrix, mask: &Matrix, scale: f32, ) -> Result<Matrix>

pub fn softplus(input: &Matrix) -> Result<Matrix>

pub fn split_heads( input: &Matrix, batch: usize, sequence_length: usize, num_heads: usize, ) -> Result<Matrix>

pub fn swiglu(gate: &Matrix, up: &Matrix) -> Result<Matrix>

pub fn tanh(input: &Matrix) -> Result<Matrix>

pub fn upsample_2d( input: &Matrix, scale_factor: usize, mode: UpsampleMode, ) -> Result<Matrix>

pub fn vq_assign(latent: &Matrix, codebook: &Matrix) -> Result<VqAssignResult>

pub fn vq_ema_update( latent: &Matrix, indices: &Matrix, embed_sum: &Matrix, cluster_size: &Matrix, codebook: &Matrix, decay: f32, epsilon: f32, dead_threshold: f32, seed: u32, normalize: bool, ) -> Result<VqEmaState>

pub fn vq_lookup(codebook: &Matrix, indices: &Matrix) -> Result<Matrix>

Function Documentation

pub fn adaptive_avg_pool_2d(
	input: &Matrix,
	output_height: usize,
	output_width: usize
) -> Result<Matrix>

Adaptively average-pool an NCHW Matrix to an exact spatial extent. Each output bin independently uses floor/ceiling boundaries along height and width. This supports rectangular and non-divisible geometry without approximating the operation as a square fixed-window pool. Errors Returns an error unless the input is nonempty rank-four F32, both output extents are nonzero, all geometry fits the shader ABI, or recording fails.

pub fn avg_pool_2d(
	input: &Matrix,
	kernel_size: usize,
	stride: usize,
	padding: usize
) -> Result<Matrix>

Apply two-dimensional average pooling to an NCHW Matrix. Padding positions do not contribute to either the sum or divisor. The output spatial extent is floor((input + 2 * padding - kernel_size) / stride) + 1, matching the OA C++ operation. Errors Returns an error unless input is a nonempty rank-four F32 Matrix, kernel_size and stride are nonzero, the kernel fits the padded input, every shader extent is representable, or runtime recording fails.

pub fn batch_norm_2d(
	input: &Matrix,
	weight: &Matrix,
	bias: &Matrix,
	epsilon: f32
) -> Result<BatchNorm2dResult>

Normalize an NCHW Matrix using statistics computed from the current batch. Errors Returns an error unless input is nonempty rank-four F32, weight and bias are same-engine F32 channel vectors, epsilon is finite and positive, or runtime recording fails.

pub fn batch_norm_2d_with_stats(
	input: &Matrix,
	mean: &Matrix,
	variance: &Matrix,
	weight: &Matrix,
	bias: &Matrix,
	epsilon: f32
) -> Result<Matrix>

Normalize an NCHW Matrix using explicit per-channel statistics. This is the stateless inference form used by crate::ml::nn::BatchNorm2d in evaluation mode. Errors Returns an error unless every value is same-engine F32, the channel-vector shapes match the input, epsilon is finite and positive, or runtime recording fails.

pub fn bmm(
	left: &Matrix,
	right: &Matrix
) -> Result<Matrix>

Multiply corresponding matrices in rank-three batches. Computes [N, M, K] @ [N, K, P] -> [N, M, P]. Errors Returns an error unless inputs are compatible nonempty same-engine F32 matrices, shape arithmetic fits the shader ABI, or runtime recording fails.

pub fn bmm_nt(
	left: &Matrix,
	right: &Matrix
) -> Result<Matrix>

Multiply rank-three batches with right-transposed storage. Computes [N, M, K] @ [N, P, K]^T -> [N, M, P]. Errors Returns an error unless inputs are compatible nonempty same-engine F32 matrices, shape arithmetic fits the shader ABI, or runtime recording fails.

pub fn bmm_tn(
	left: &Matrix,
	right: &Matrix
) -> Result<Matrix>

Multiply rank-three batches with left-transposed storage. Computes [N, K, M]^T @ [N, K, P] -> [N, M, P]. Errors Returns an error unless inputs are compatible nonempty same-engine F32 matrices, shape arithmetic fits the shader ABI, or runtime recording fails.

pub fn channel_norm(
	input: &Matrix,
	weight: &Matrix,
	bias: &Matrix,
	epsilon: f32
) -> Result<Matrix>

Normalize the channel axis of a nonempty FP32 [B, C, T] Matrix. Errors Returns an error unless the affine vectors match C, all values belong to one Engine, epsilon is finite and positive, C <= 1024, or recording fails.

pub fn channel_norm_backward(
	input: &Matrix,
	weight: &Matrix,
	output_gradient: &Matrix,
	epsilon: f32
) -> Result<ChannelNormBackward>

Compute complete explicit channel-normalization adjoints. Errors Returns an error when the forward contract or output gradient is invalid, or runtime recording fails.

pub fn channel_norm_relu(
	input: &Matrix,
	weight: &Matrix,
	bias: &Matrix,
	epsilon: f32
) -> Result<Matrix>

Normalize the channel axis of [B, C, T] and fuse the following ReLU. Errors Returns an error under the same conditions as channel_norm, or when runtime recording fails.

pub fn channel_norm_relu_backward(
	input: &Matrix,
	weight: &Matrix,
	forward_output: &Matrix,
	output_gradient: &Matrix,
	epsilon: f32
) -> Result<ChannelNormBackward>

Compute complete adjoints of fused channel normalization and ReLU. Errors Returns an error when saved output or gradient does not match the forward contract, or runtime recording fails.

pub fn conv_1d(
	input: &Matrix,
	weight: &Matrix,
	bias: &Matrix,
	stride: usize,
	padding: usize,
	dilation: usize
) -> Result<Matrix>

Apply one-dimensional convolution to an NCL Matrix through im2col and GEMM. weight uses [output_channels, input_channels, kernel] layout and bias uses [output_channels]. Dilation applies within each kernel window. Errors Returns an error unless all inputs are nonempty same-engine F32 matrices, shapes agree, stride and dilation are nonzero, the effective kernel fits the padded input, every extent fits the shader ABI, or recording fails.

pub fn conv_2d(
	input: &Matrix,
	weight: &Matrix,
	bias: &Matrix,
	stride: usize,
	padding: usize,
	groups: usize
) -> Result<Matrix>

Apply grouped two-dimensional convolution to an NCHW Matrix. weight uses [output_channels, input_channels / groups, kernel, kernel] layout and bias uses [output_channels]. The output spatial extent is floor((input + 2 * padding - kernel) / stride) + 1. Errors Returns an error unless all inputs are same-engine F32 matrices, shapes and group divisibility agree, the square kernel fits the padded input, every extent fits the shader ABI, or runtime recording fails.

pub fn conv_transpose_1d(
	input: &Matrix,
	weight: &Matrix,
	stride: usize,
	padding: usize,
	dilation: usize
) -> Result<Matrix>

Apply bias-free one-dimensional transposed convolution to an NCL Matrix. weight uses [input_channels, output_channels, kernel] layout. The output length is (input_length - 1) * stride - 2 * padding + effective_kernel. Errors Returns an error unless both inputs are nonempty same-engine F32 matrices, channel geometry agrees, stride and dilation are nonzero, the output extent is positive and fits the shader ABI, or runtime recording fails.

pub fn conv_transpose_2d(
	input: &Matrix,
	weight: &Matrix,
	bias: &Matrix,
	stride: usize,
	padding: usize
) -> Result<Matrix>

Apply two-dimensional transposed convolution to an NCHW Matrix. weight uses [input_channels, output_channels, kernel, kernel] layout and bias uses [output_channels]. Each output spatial extent is (input_extent - 1) * stride - 2 * padding + kernel. Errors Returns an error unless all inputs are nonempty same-engine F32 matrices, channel and square-kernel geometry agree, stride is nonzero, output extents are positive and fit the shader ABI, or runtime recording fails.

pub fn detach(
	input: &Matrix
) -> Result<Matrix>

Return a zero-copy stop-gradient view. The returned Matrix shares storage and forward values with input, but owns a fresh semantic identity with no reverse-mode edge to input. Errors Returns an error only if the input’s existing shape metadata is invalid.

pub fn elu(
	input: &Matrix,
	alpha: f32
) -> Result<Matrix>

Apply the exponential linear unit elementwise. Errors Returns an error unless input is F32 or runtime recording fails.

pub fn flash_attention_causal(
	query: &Matrix,
	key: &Matrix,
	value: &Matrix,
	scale: f32
) -> Result<Matrix>

Apply the explicit causal Flash provider to [BH,S,Dh] values. This compatibility entry point retains the same semantic scaled-dot-product attention identity as scaled_dot_product_attention. It materializes only the output and one FP32 log-sum-exp value per query row. Errors Returns an error unless Q, K, and V are equal nonempty same-engine F32 matrices, the sequence length is at most 1024, or runtime recording fails.

pub fn gelu(
	input: &Matrix
) -> Result<Matrix>

Apply the tanh-approximate Gaussian error linear unit elementwise. Errors Returns an error unless input is F32 or runtime recording fails.

pub fn grouped_gemm_m(
	input: &Matrix,
	weight: &Matrix,
	offsets: &Matrix
) -> Result<Matrix>

Apply one bias-free projection per expert to expert-major packed rows. input is F32 [R, K], weight is F32 [E, N, K], and offsets is the U32 [E+1] boundary output from crate::matrix::moe_expert_plan. The result is F32 [R, N]. Errors Returns an error for incompatible shapes, dtypes, engine ownership, shader ABI overflow, allocation failure, or runtime recording failure.

pub fn grouped_linear_m(
	input: &Matrix,
	weight: &Matrix,
	bias: &Matrix,
	offsets: &Matrix
) -> Result<Matrix>

Apply one affine projection per expert to expert-major packed rows. input is F32 [R, K], weight is F32 [E, N, K], bias is F32 [E, N], and offsets is the U32 [E+1] boundary output from crate::matrix::moe_expert_plan. The result is F32 [R, N]. Errors Returns an error for incompatible shapes, dtypes, engine ownership, shader ABI overflow, allocation failure, or runtime recording failure.

pub fn gru_cell(
	gates_i: &Matrix,
	hidden: &Matrix,
	weight_hh: &Matrix,
	bias_hh: Option<&Matrix>
) -> Result<Matrix>

Apply one GRU cell update to a caller-supplied hidden state. gates_i has shape [batch, 3 * hidden], hidden has [batch, hidden], weight_hh has [3 * hidden, hidden], and an optional recurrent bias has [3 * hidden]. Errors Returns an error unless all values are nonempty same-engine F32 matrices, shapes agree, hidden size is at most 1024, or runtime recording fails.

pub fn gru_scan(
	gates_i: &Matrix,
	weight_hh: &Matrix,
	bias_hh: Option<&Matrix>
) -> Result<Matrix>

Run a whole-sequence GRU recurrence from a zero hidden state. gates_i has shape [batch, sequence, 3 * hidden], weight_hh has [3 * hidden, hidden], and an optional bias_hh has [3 * hidden]. Gate order is reset, update, candidate. Errors Returns an error unless all values are nonempty same-engine F32 matrices, shapes agree, hidden size is at most 1024, or runtime recording fails.

pub fn leaky_relu(
	input: &Matrix,
	alpha: f32
) -> Result<Matrix>

Apply Leaky ReLU elementwise with the requested negative slope. Errors Returns an error unless input is F32 or runtime recording fails.

pub fn mamba3_mimo(
	c: &Matrix,
	b: &Matrix,
	x: &Matrix,
	z: &Matrix,
	adt: &Matrix,
	dt: &Matrix,
	trap: &Matrix,
	angle: &Matrix,
	c_bias: &Matrix,
	b_bias: &Matrix,
	d: &Matrix,
	mimo_x: &Matrix,
	mimo_z: &Matrix,
	mimo_o: &Matrix,
	norm_weight: &Matrix,
	config: SsmConfig
) -> Result<Matrix>

Run the shared-state Mamba-3 MIMO selective scan. C/B use [B,L,R*G,N]; the three MIMO projections use [H,R,P]; the output-normalization weight always uses [H,P] and is consumed only when config.has_output_norm is true. The remaining layouts match mamba3_siso. Errors Returns an error for mismatched geometry, rank outside 1..=8, unsupported dimensions, shape overflow, or failed runtime recording.

pub fn mamba3_mimo_backward(
	output_gradient: &Matrix,
	c: &Matrix,
	b: &Matrix,
	x: &Matrix,
	z: &Matrix,
	adt: &Matrix,
	dt: &Matrix,
	trap: &Matrix,
	angle: &Matrix,
	c_bias: &Matrix,
	b_bias: &Matrix,
	d: &Matrix,
	mimo_x: &Matrix,
	mimo_z: &Matrix,
	mimo_o: &Matrix,
	norm_weight: &Matrix,
	config: SsmConfig
) -> Result<Mamba3MimoBackward>

Compute all fifteen explicit Mamba-3 MIMO adjoints. The implementation records the donor’s deterministic prepare, post, reverse-core, inverse-rotation, and reduction stages as one semantic operation. Errors Returns an error for mismatched geometry, unsupported dimensions, shape overflow, or failed runtime recording.

pub fn mamba3_mimo_step(
	c: &Matrix,
	b: &Matrix,
	x: &Matrix,
	z: &Matrix,
	adt: &Matrix,
	dt: &Matrix,
	trap: &Matrix,
	angle: &Matrix,
	c_bias: &Matrix,
	b_bias: &Matrix,
	d: &Matrix,
	mimo_x: &Matrix,
	mimo_z: &Matrix,
	mimo_o: &Matrix,
	norm_weight: &Matrix,
	ssm_state: &mut Matrix,
	angle_state: &mut Matrix,
	k_state: &mut Matrix,
	v_state: &mut Matrix,
	config: SsmConfig
) -> Result<Matrix>

Advance one Mamba-3 MIMO token and update four explicit recurrent states. State layouts are SSM [B,H,P,N], angle [B,H,A], key [B,H,R,N], and value [B,H,R,P]. Errors Returns an error unless config.sequence_length == 1, every operand and state matches the configured same-engine F32 layout, or recording fails.

pub fn mamba3_preprocess(
	projected: &Matrix,
	dt_bias: &Matrix,
	config: Mamba3PreprocessConfig
) -> Result<Mamba3PreprocessResult>

Split and transform one packed Mamba-3 input projection. projected has shape [rows, projected_width], where `projected_width = 2inner_size + 2num_groupsmimo_rankstate_size 3*num_heads + num_rope_angles. dt_bias` has one element per head. The operation records one semantic node and one physical dispatch. Errors Returns an error for invalid configuration, incompatible shape, dtype or Engine ownership, arithmetic overflow, or failed runtime recording.

pub fn mamba3_preprocess_backward(
	projected: &Matrix,
	dt_bias: &Matrix,
	output_gradients: &Mamba3PreprocessResult,
	config: Mamba3PreprocessConfig
) -> Result<Mamba3PreprocessBackward>

Compute the packed projection and dt-bias adjoints for Mamba-3 preprocess. output_gradients follows the exact eight shapes returned by mamba3_preprocess. Its dt-bias reduction is ordered by increasing row. Errors Returns an error for invalid configuration, incompatible gradients, shape, dtype or Engine ownership, arithmetic overflow, or failed runtime recording.

pub fn mamba3_siso(
	c: &Matrix,
	b: &Matrix,
	x: &Matrix,
	z: &Matrix,
	adt: &Matrix,
	dt: &Matrix,
	trap: &Matrix,
	angle: &Matrix,
	c_bias: &Matrix,
	b_bias: &Matrix,
	d: &Matrix,
	config: SsmConfig
) -> Result<Matrix>

Run the Mamba-3 SISO selective-state recurrence. Operand layouts are C/B [B,L,G,N], x/z/output [B,L,H,P], A-times-dt/dt/trap [B,L,H], angle [B,L,A], C/B bias [H,N], and D [H]. All operands are same-engine F32 matrices. Errors Returns an error for mismatched geometry, unsupported dimensions, shape overflow, or failed runtime recording.

pub fn mamba3_siso_backward(
	output_gradient: &Matrix,
	c: &Matrix,
	b: &Matrix,
	x: &Matrix,
	z: &Matrix,
	adt: &Matrix,
	dt: &Matrix,
	trap: &Matrix,
	angle: &Matrix,
	c_bias: &Matrix,
	b_bias: &Matrix,
	d: &Matrix,
	config: SsmConfig
) -> Result<SsmBackward>

Compute all explicit Mamba-3 SISO adjoints. The initial OARS physical route is the donor short-training route and accepts L <= 16, P <= 16, N <= 32, and A <= 8. Errors Returns an error for mismatched inputs, a shape outside the admitted short route, overflow, or failed runtime recording.

pub fn mamba3_siso_step(
	c: &Matrix,
	b: &Matrix,
	x: &Matrix,
	z: &Matrix,
	adt: &Matrix,
	dt: &Matrix,
	trap: &Matrix,
	angle: &Matrix,
	c_bias: &Matrix,
	b_bias: &Matrix,
	d: &Matrix,
	ssm_state: &mut Matrix,
	angle_state: &mut Matrix,
	k_state: &mut Matrix,
	v_state: &mut Matrix,
	config: SsmConfig
) -> Result<Matrix>

Advance one Mamba-3 SISO token and update four explicit recurrent states. Operand layouts are the length-one forms accepted by mamba3_siso. State layouts are SSM [B,H,P,N], cumulative angle [B,H,A], previous key [B,H,N], and previous value [B,H,P]. The states are mutated only after the complete operation has been validated and recorded. Errors Returns an error unless config.sequence_length == 1, every operand and state has the configured same-engine F32 shape, or recording fails.

pub fn max_pool_2d(
	input: &Matrix,
	kernel_size: usize,
	stride: usize,
	padding: usize
) -> Result<MaxPool2dResult>

Apply two-dimensional max pooling to an NCHW Matrix. Equal maxima select the first valid input in row-major window order. The returned U32 argmax indices are explicit because they are part of the donor operation contract and make the deterministic adjoint reusable. Errors Returns an error under the same geometry and runtime conditions as avg_pool_2d.

pub fn merge_heads(
	input: &Matrix,
	batch: usize,
	sequence_length: usize,
	num_heads: usize
) -> Result<Matrix>

Invert split_heads into contiguous [B*S, H*(D/H)] storage. The one-head case is a zero-copy differentiable view. Errors Returns an error unless input is nonempty F32 [B*H, S, D/H] with the supplied nonzero dimensions, or runtime recording fails.

pub fn mish(
	input: &Matrix
) -> Result<Matrix>

Apply Mish elementwise. Errors Returns an error unless input is F32 or runtime recording fails.

pub fn moe_combine(
	packed: &Matrix,
	route_gate: &Matrix,
	inverse: &Matrix,
	packed_slot: &Matrix
) -> Result<Matrix>

Combine expert-major packed rows into token-major output using route gates. packed is F32 [T*K, D], route_gate is F32 [T, K], and the U32 maps come from crate::matrix::moe_expert_plan. Errors Returns an error for incompatible shapes, dtypes, engine ownership, shader ABI overflow, allocation failure, or runtime recording failure.

pub fn moe_gather(
	input: &Matrix,
	packed_token: &Matrix,
	inverse: &Matrix
) -> Result<Matrix>

Pack token rows into stable expert-major route order. input is F32 [T, D]; packed_token and inverse are the U32 [T*K] maps produced by crate::matrix::moe_expert_plan. inverse is retained for the deterministic reverse pass and is not read by the forward kernel. Errors Returns an error for incompatible shapes, dtypes, engine ownership, shader ABI overflow, allocation failure, or runtime recording failure.

pub fn moe_route_weights(
	probabilities: &Matrix,
	expert_indices: &Matrix
) -> Result<Matrix>

Normalize probabilities over the selected experts for each token. probabilities is F32 [T, E]; expert_indices is I32 [T, K] with 1 <= K <= E. The routing decision is detached from reverse mode, while gradients propagate to the selected probability entries. Errors Returns an error for incompatible shapes, dtypes, engine ownership, shader ABI overflow, allocation failure, or runtime recording failure.

pub fn relu(
	input: &Matrix
) -> Result<Matrix>

Apply the rectified linear unit elementwise. Errors Returns an error unless input is F32 or runtime recording fails.

pub fn rms_norm(
	input: &Matrix,
	weight: &Matrix,
	epsilon: f32
) -> Result<Matrix>

Apply weighted root-mean-square normalization over the final dimension. Errors Returns an error unless input is a nonempty F32 Matrix, weight is a same-engine F32 vector matching its final dimension, epsilon is finite and positive, or runtime recording fails.

pub fn rms_norm_gated(
	input: &Matrix,
	weight: &Matrix,
	bias: Option<&Matrix>,
	gate: &Matrix,
	epsilon: f32
) -> Result<Matrix>

Apply per-row RMS normalization, a cyclic broadcast affine, and a SiLU gate. Normalization is over the final input dimension. The affine value may have one or more leading groups; those groups repeat over the flattened input rows. An optional bias must exactly match weight. Errors Returns an error for incompatible shape, dtype, Engine ownership or epsilon, or when runtime recording fails.

pub fn rms_norm_gated_backward(
	input: &Matrix,
	weight: &Matrix,
	bias: Option<&Matrix>,
	gate: &Matrix,
	output_gradient: &Matrix,
	epsilon: f32
) -> Result<RmsNormGatedBackward>

Compute all gated RMS normalization adjoints explicitly. Errors Returns an error for an incompatible output gradient or forward contract, or when runtime recording fails.

pub fn rnn_cell(
	gates_i: &Matrix,
	hidden: &Matrix,
	weight_hh: &Matrix,
	bias_hh: Option<&Matrix>
) -> Result<Matrix>

Apply one Elman RNN cell update to a caller-supplied hidden state. gates_i and hidden have shape [batch, hidden], weight_hh has [hidden, hidden], and an optional recurrent bias has [hidden]. Errors Returns an error unless all values are nonempty same-engine F32 matrices, shapes agree, hidden size is at most 1024, or runtime recording fails.

pub fn rnn_scan(
	gates_i: &Matrix,
	weight_hh: &Matrix,
	bias_hh: Option<&Matrix>
) -> Result<Matrix>

Run a whole-sequence Elman recurrence from a zero hidden state. gates_i has shape [batch, sequence, hidden], weight_hh has [hidden, hidden], and an optional recurrent bias has [hidden]. Errors Returns an error unless all values are nonempty same-engine F32 matrices, shapes agree, hidden size is at most 1024, or runtime recording fails.

pub fn rope(
	input: &Matrix,
	num_heads: usize,
	head_dim: usize,
	theta_base: f32,
	position_offset: u32
) -> Result<Matrix>

Apply split-half rotary position embeddings to [tokens, heads * head_dim]. position_offset selects the first absolute position, making the same operation usable for full sequences and future KV-cache extension. Errors Returns an error unless the input is nonempty F32, head_dim is even, the final width matches num_heads * head_dim, and theta_base is finite and positive.

pub fn scaled_dot_product_attention(
	query: &Matrix,
	key: &Matrix,
	value: &Matrix,
	additive_mask: Option<&Matrix>,
	scale: f32,
	causal: bool
) -> Result<Matrix>

Apply standard scaled dot-product attention to [B*H, S, D/H] values. additive_mask, when present, is detached F32 [B*H*S, S] storage. Causal visibility is applied without materializing another mask. Errors Returns an error unless Q, K, and V are equal nonempty same-engine F32 rank-three matrices, the optional mask matches their flattened score rows, or runtime recording fails.

pub fn scaled_dot_product_attention_causal(
	query: &Matrix,
	key: &Matrix,
	value: &Matrix,
	sequence_length: usize,
	num_heads: usize
) -> Result<Matrix>

Apply multi-head causal scaled dot-product attention to packed [B*S, D] values. Errors Returns an error unless Q, K, and V are equal nonempty same-engine F32 matrices, their row count is divisible by sequence_length, their width is divisible by num_heads, or runtime recording fails.

pub fn sigmoid(
	input: &Matrix
) -> Result<Matrix>

Apply the logistic sigmoid elementwise. Errors Returns an error unless input is F32 or runtime recording fails.

pub fn silu(
	input: &Matrix
) -> Result<Matrix>

Apply the sigmoid linear unit elementwise. Errors Returns an error unless input is F32 or runtime recording fails.

pub fn silu_mul(
	input: &Matrix,
	intermediate_size: usize
) -> Result<Matrix>

Apply SiLU(gate) * up to concatenated halves of the final input axis. Errors Returns an error unless input is F32 with final extent exactly twice the nonzero intermediate_size, or runtime recording fails.

pub fn softmax_scaled_masked(
	scores: &Matrix,
	mask: &Matrix,
	scale: f32
) -> Result<Matrix>

Compute softmax(scores * scale + mask) over the final axis. The additive mask is detached from reverse mode. Errors Returns an error unless both values are equal-shape, nonempty, same-engine F32 matrices, or runtime recording fails.

pub fn softplus(
	input: &Matrix
) -> Result<Matrix>

Apply the overflow-safe Softplus formulation elementwise. Errors Returns an error unless input is F32 or runtime recording fails.

pub fn split_heads(
	input: &Matrix,
	batch: usize,
	sequence_length: usize,
	num_heads: usize
) -> Result<Matrix>

Permute contiguous [B*S, D] storage into [B*H, S, D/H] heads. The one-head case is a zero-copy differentiable view. Errors Returns an error unless input is nonempty F32 [B*S, D], all dimensions are nonzero, D is divisible by num_heads, or runtime recording fails.

pub fn swiglu(
	gate: &Matrix,
	up: &Matrix
) -> Result<Matrix>

Apply SiLU(gate) * up elementwise. Errors Returns an error unless both inputs are equal-shape F32 on the same engine, or runtime recording fails.

pub fn tanh(
	input: &Matrix
) -> Result<Matrix>

Apply hyperbolic tangent elementwise. Errors Returns an error unless input is F32 or runtime recording fails.

pub fn upsample_2d(
	input: &Matrix,
	scale_factor: usize,
	mode: UpsampleMode
) -> Result<Matrix>

Upsample an FP32 NCHW Matrix by a positive integer scale factor. Errors Returns an error unless the input is nonempty rank-four F32, the scale factor is positive, all output geometry is representable, or runtime recording fails.

pub fn vq_assign(
	latent: &Matrix,
	codebook: &Matrix
) -> Result<VqAssignResult>

Assign every FP32 latent row to its nearest FP32 codebook row. Squared L2 distance is used and equal-distance ties select the lower code index, matching OA C++. The hard assignment is detached from reverse mode. Errors Returns an error unless inputs are nonempty same-engine FP32 [N,D] and [K,D] matrices or recording/allocation fails.

pub fn vq_ema_update(
	latent: &Matrix,
	indices: &Matrix,
	embed_sum: &Matrix,
	cluster_size: &Matrix,
	codebook: &Matrix,
	decay: f32,
	epsilon: f32,
	dead_threshold: f32,
	seed: u32,
	normalize: bool
) -> Result<VqEmaState>

Produce the next VQ EMA state without mutating input Matrix values. This preserves the donor update equations, optional unit-RMS codebook normalization, and exact dead-code revival hash. The owning module installs the returned values into its stable persistent buffer handles. Errors Returns an error for incompatible shapes, dtypes, engines, invalid numerical options, ABI overflow, allocation failure, or runtime recording failure.

pub fn vq_lookup(
	codebook: &Matrix,
	indices: &Matrix
) -> Result<Matrix>

Gather codebook rows selected by I32 token indices. Invalid indices yield zero rows without an out-of-bounds read. This pure decode operation is detached from reverse mode because EMA owns codebook learning. Errors Returns an error unless the codebook is FP32 [K,D], indices are I32 [N], both belong to one engine, or recording/allocation fails.