oa::matrix

matrix module in the OA Core public surface.

MdRust public API

Public Types

Public Functions

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

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

pub fn add_scalar(input: &Matrix, scalar: f32) -> Result<Matrix>

pub fn categorical_accuracy_count( logits: &Matrix, labels: &Matrix, ) -> Result<Matrix>

pub fn clamp_max(input: &Matrix, maximum: f32) -> Result<Matrix>

pub fn clamp_min(input: &Matrix, minimum: f32) -> Result<Matrix>

pub fn concat(inputs: &[Matrix], dim: i32) -> Result<Matrix>

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

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

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

pub fn div_scalar(input: &Matrix, scalar: f32) -> Result<Matrix>

pub fn dropout(input: &Matrix, probability: f32, seed: u64) -> Result<Matrix>

pub fn equal(input: &Matrix, value: f32) -> Result<Matrix>

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

pub fn from_f32(shape: impl Into<Vec<usize>>, values: &[f32]) -> Result<Matrix>

pub fn from_f32_on( engine: &Engine, shape: impl Into<Vec<usize>>, values: &[f32], ) -> Result<Matrix>

pub fn full( engine: &Engine, shape: impl Into<Vec<usize>>, value: f32, ) -> Result<Matrix>

pub fn gather(input: &Matrix, indices: &Matrix) -> Result<Matrix>

pub fn gather_last_dim(input: &Matrix, indices: &Matrix) -> Result<Matrix>

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

pub fn log_softmax(input: &Matrix, dim: i32) -> Result<Matrix>

pub fn masked_categorical_accuracy_count( logits: &Matrix, labels: &Matrix, mask: &Matrix, ) -> Result<Matrix>

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

pub fn moe_expert_plan( indices: &Matrix, num_experts: usize, ) -> Result<MoeExpertPlan>

pub fn moe_routing_bias_update( selection_mask: &Matrix, bias: &Matrix, experts_per_token: usize, gamma: f32, ) -> Result<()>

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

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

pub fn ones(engine: &Engine, shape: impl Into<Vec<usize>>) -> Result<Matrix>

pub fn philox_normal( input: &Matrix, mean: f32, stddev: f32, seed: u64, ) -> Result<Matrix>

pub fn philox_uniform( input: &Matrix, low: f32, high: f32, seed: u64, ) -> Result<Matrix>

pub fn pow(input: &Matrix, exponent: f32) -> Result<Matrix>

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

pub fn repeat_interleave( input: &Matrix, repeats: usize, dim: i32, ) -> Result<Matrix>

pub fn reshape(input: &Matrix, shape: impl Into<Vec<usize>>) -> Result<Matrix>

pub fn sample_logits( logits: &Matrix, temperature: f32, top_k: i32, top_p: f32, seed: u64, ) -> Result<Matrix>

pub fn scale(input: &Matrix, scalar: f32) -> Result<Matrix>

pub fn set_rng_seed(seed: u64)

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

pub fn slice(input: &Matrix, dim: i32, start: i64, end: i64) -> Result<Matrix>

pub fn softmax(input: &Matrix, dim: i32) -> Result<Matrix>

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

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

pub fn sub_scalar(input: &Matrix, scalar: f32) -> Result<Matrix>

pub fn sum(input: &Matrix, dim: i32) -> Result<Matrix>

pub fn top_k(input: &Matrix, k: i32, dim: i32) -> Result<TopKResult>

pub fn top_k_mask(indices: &Matrix, num_experts: usize) -> Result<Matrix>

pub fn transpose(input: &Matrix, dim0: i32, dim1: i32) -> Result<Matrix>

Function Documentation

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

Take the absolute value of every FP32 matrix element. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

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

Add two matrices elementwise using multidirectional broadcasting. Shapes use multidirectional broadcasting over at most eight axes. The I32 route uses two’s-complement wrapping arithmetic. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn add_scalar(
	input: &Matrix,
	scalar: f32
) -> Result<Matrix>

Add a scalar to every FP32 matrix element. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn categorical_accuracy_count(
	logits: &Matrix,
	labels: &Matrix
) -> Result<Matrix>

Count rows whose last-axis FP32 argmax equals the integer class label. The returned [1] U32 Matrix remains device-resident. Reading it is an explicit synchronization boundary. Errors Returns an error when logits are not a nonempty rank-two-or-greater FP32 Matrix, labels do not match the leading logits shape or use U8/U32/I32, inputs belong to different engines, a size exceeds the shader ABI, allocation fails, or recording fails.

pub fn clamp_max(
	input: &Matrix,
	maximum: f32
) -> Result<Matrix>

Cap every FP32 matrix element at an upper bound. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn clamp_min(
	input: &Matrix,
	minimum: f32
) -> Result<Matrix>

Raise every FP32 matrix element to a lower bound. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn concat(
	inputs: &[Matrix],
	dim: i32
) -> Result<Matrix>

Concatenate one or more same-dtype matrices along one axis. This is the direct Rust spelling of OA C++ FnMatrix::concat. One semantic operation lowers to one copy-region dispatch per nonempty input. Errors Returns an error when there are no inputs, rank is outside one through four, dtype, engine, rank, or non-concatenated extents differ, the axis is invalid, shape or shader-ABI arithmetic overflows, allocation fails, or recording fails.

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

Deep-copy a dense matrix into independent device storage. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

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

Take the cosine of every FP32 matrix element. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

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

Divide two FP32 matrices elementwise using multidirectional broadcasting. Shapes use multidirectional broadcasting over at most eight axes. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn div_scalar(
	input: &Matrix,
	scalar: f32
) -> Result<Matrix>

Divide every FP32 matrix element by a scalar. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn dropout(
	input: &Matrix,
	probability: f32,
	seed: u64
) -> Result<Matrix>

Apply inverted Dropout with a Philox-generated mask. Kept values are scaled by 1 / (1 - probability). The effective seed is retained for the reverse pass, and captured execution advances its random counter once per replay without rebuilding the graph. Errors Returns an error unless input is FP32, probability is finite and in [0, 1), or runtime allocation/recording fails.

pub fn equal(
	input: &Matrix,
	value: f32
) -> Result<Matrix>

Compare every FP32 element with one scalar and return an FP32 zero/one mask. Equality follows the device’s ordinary IEEE comparison: NaN never compares equal, while positive and negative zero compare equal. The result is intentionally detached because a discrete mask has no admitted adjoint. Errors Returns an error when input is not FP32, its element count exceeds the shader ABI, allocation fails, or runtime recording fails.

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

Take the exponential of every FP32 matrix element. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn from_f32(
	shape: impl Into<Vec<usize>>,
	values: &[f32]
) -> Result<Matrix>

Create an FP32 matrix on this thread’s default engine. Errors Returns an error when the default engine cannot be initialized, the shape does not match values, or allocation and upload fail.

pub fn from_f32_on(
	engine: &Engine,
	shape: impl Into<Vec<usize>>,
	values: &[f32]
) -> Result<Matrix>

Create an FP32 matrix on a particular engine. Errors Returns an error when the shape does not match values, or allocation and upload fail.

pub fn full(
	engine: &Engine,
	shape: impl Into<Vec<usize>>,
	value: f32
) -> Result<Matrix>

Create an FP32 matrix filled with value. Errors Returns an error when the shape or storage size overflows, or allocation and upload fail.

pub fn gather(
	input: &Matrix,
	indices: &Matrix
) -> Result<Matrix>

Gather rows from an FP32 table using arbitrary-shape integer indices. The table must be [rows,width]; indices are flattened and output is [num_indices,width], matching OA C++ FnMatrix::gather. U8, U32, and I32 indices use one bounded kernel. Invalid rows produce NaN without reading outside the table. Errors Returns an error for incompatible rank, dtype, Engine ownership, shader ABI, allocation, or recording.

pub fn gather_last_dim(
	input: &Matrix,
	indices: &Matrix
) -> Result<Matrix>

Gather I32-selected elements from every row’s final axis. input is FP32 [R, C] and indices is I32 [R, K]; the result is FP32 [R, K]. Out-of-range indices produce zero, matching OA C++. Repeated indices accumulate deterministically in the reverse pass. Errors Returns an error unless both inputs are rank two on the same engine, row counts agree, dtypes match the contract, the input width is positive, sizes fit the shader ABI, allocation fails, or runtime recording fails.

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

Take the natural logarithm of every FP32 matrix element. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn log_softmax(
	input: &Matrix,
	dim: i32
) -> Result<Matrix>

Stable LogSoftmax over dim, where -1 selects the last dimension. The operation records asynchronously and returns its output immediately. Errors Returns an error when the input is empty or not FP32, dim is invalid, an axis extent or dispatch count exceeds the shader ABI, allocation fails, or runtime recording fails.

pub fn masked_categorical_accuracy_count(
	logits: &Matrix,
	labels: &Matrix,
	mask: &Matrix
) -> Result<Matrix>

Count correct categorical predictions only where the FP32 mask is nonzero. The returned [1] U32 Matrix remains device-resident. Reading it is an explicit synchronization boundary. Errors Returns an error under the same conditions as categorical_accuracy_count, or when the mask is not FP32, does not match the labels shape, or belongs to another engine.

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

Multiply A by the transpose of B using the OA weight-layout convention. left has shape [M, K], right has shape [N, K], and the returned matrix has shape [M, N]. This matches OA’s native weight-layout convention and does not transpose storage. The call records device work and returns without a host wait. Errors Returns an error when either input is not a rank-two FP32 matrix, their K extents differ, they belong to different engines, a dimension exceeds the admitted shader ABI, or allocation or recording fails.

pub fn moe_expert_plan(
	indices: &Matrix,
	num_experts: usize
) -> Result<MoeExpertPlan>

Build a stable dropless expert-major plan from I32 [T, K] indices. Within every expert, routes preserve their original token-major order. The operation accepts at most 256 experts, matching the shader’s shared state. Errors Returns an error unless indices are nonempty rank-two I32, the expert count is in 1..=256, sizes fit the admitted ABI, allocation succeeds, or runtime recording succeeds.

pub fn moe_routing_bias_update( selection_mask: &Matrix, bias: &Matrix, experts_per_token: usize, gamma: f32, ) -> Result<()>

Update an FP32 expert-routing bias from a dense FP32 selection mask. selection_mask is [T, E]; bias may be [E] or [1, E]. Each expert receives -gamma, +gamma, or no update when its observed route count is above, below, or equal to the balanced target T * K / E. The update stays device-resident and mutates bias through its existing value identity. Errors Returns an error unless both matrices are FP32 on the same engine, the mask and bias shapes agree, experts_per_token is in 1..=E, gamma is finite and positive, dimensions fit the shader ABI, or runtime recording fails.

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

Multiply two FP32 matrices elementwise. Shapes use multidirectional broadcasting over at most eight axes. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

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

Negate every FP32 matrix element. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn ones(
	engine: &Engine,
	shape: impl Into<Vec<usize>>
) -> Result<Matrix>

Create an FP32 matrix filled with ones. Errors Returns an error when the shape or storage size overflows, or allocation and upload fail.

pub fn philox_normal(
	input: &Matrix,
	mean: f32,
	stddev: f32,
	seed: u64
) -> Result<Matrix>

Generate normally distributed FP32 values with input’s shape. Errors Returns an error unless input is FP32, mean and stddev are finite, stddev is nonnegative, or runtime allocation/recording fails.

pub fn philox_uniform(
	input: &Matrix,
	low: f32,
	high: f32,
	seed: u64
) -> Result<Matrix>

Generate uniformly distributed FP32 values with input’s shape. input supplies shape, dtype, engine, and semantic lineage; its values are not read. A nonzero seed is exact. Zero consumes the thread-local seed sequence controlled by set_rng_seed. Errors Returns an error unless input is FP32, low < high, both bounds are finite, or runtime allocation/recording fails.

pub fn pow(
	input: &Matrix,
	exponent: f32
) -> Result<Matrix>

Raise every FP32 matrix element to a scalar exponent. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

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

Take the reciprocal of every FP32 matrix element. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn repeat_interleave(
	input: &Matrix,
	repeats: usize,
	dim: i32
) -> Result<Matrix>

Repeat every FP32 element along one matrix axis. This is the direct Rust spelling of OA C++ FnMatrix::repeatInterleave. Rank-one through rank-four matrices are currently admitted. Errors Returns an error for an unsupported rank or dtype, non-positive repeat count, invalid axis, shape/ABI overflow, allocation failure, or recording failure.

pub fn reshape(
	input: &Matrix,
	shape: impl Into<Vec<usize>>
) -> Result<Matrix>

Return a zero-copy dense view with a different shape. The result owns a new semantic value identity but shares storage, pending execution state, dtype, and element order with input. When a gradient tape is active, its adjoint reshapes the incoming gradient back to the input shape. Errors Returns an error when the requested shape overflows or changes the number of elements.

pub fn sample_logits(
	logits: &Matrix,
	temperature: f32,
	top_k: i32,
	top_p: f32,
	seed: u64
) -> Result<Matrix>

Sample one I32 class index from each final-axis row of FP32 logits. Nonpositive temperature selects the first maximum deterministically. Positive temperature uses the donor Philox categorical route; positive top_k restricts candidates and top_p applies nucleus truncation after descending TopK. Rank-one input is treated as one row. Errors Returns an error when logits are not nonempty rank-one or rank-two FP32, a sorted candidate count exceeds 1024, sizes do not fit the shader ABI, allocation fails, or runtime recording fails.

pub fn scale(
	input: &Matrix,
	scalar: f32
) -> Result<Matrix>

Multiply every FP32 matrix element by a scalar. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn set_rng_seed(
	seed: u64
)

Replace the current thread’s implicit random seed source. Explicit nonzero operation seeds bypass this state. Passing zero to a Philox operation consumes one value from this deterministic thread-local sequence.

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

Take the sine of every FP32 matrix element. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn slice(
	input: &Matrix,
	dim: i32,
	start: i64,
	end: i64
) -> Result<Matrix>

Materialize a half-open interval along one matrix axis. This is the direct Rust spelling of OA C++ FnMatrix::slice. The current checkpoint admits rank-one through rank-four FP32 matrices. Errors Returns an error when the rank, dtype, axis, range, shader ABI, allocation, or runtime recording contract is invalid.

pub fn softmax(
	input: &Matrix,
	dim: i32
) -> Result<Matrix>

Stable Softmax over dim, where -1 selects the last dimension. The operation records asynchronously and returns its output immediately. Errors Returns an error when the input is empty or not FP32, dim is invalid, an axis extent or dispatch count exceeds the shader ABI, allocation fails, or runtime recording fails.

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

Take the square root of every FP32 matrix element. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

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

Subtract two FP32 matrices elementwise using multidirectional broadcasting. Shapes use multidirectional broadcasting over at most eight axes. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn sub_scalar(
	input: &Matrix,
	scalar: f32
) -> Result<Matrix>

Subtract a scalar from every FP32 matrix element. The call records device work and returns without a host wait. Errors Returns an error when the shape, dtype, ownership, allocation, or recording contract fails.

pub fn sum(
	input: &Matrix,
	dim: i32
) -> Result<Matrix>

Sum all values when dim is -1, or keep and reduce one selected axis. A full reduction returns shape [1]; an axis reduction preserves rank and replaces the selected extent with one. The operation records asynchronously. Errors Returns an error when the input is empty or not FP32, dim is invalid, a size exceeds the admitted shader ABI, allocation fails, or recording fails.

pub fn top_k(
	input: &Matrix,
	k: i32,
	dim: i32
) -> Result<TopKResult>

Return the largest k FP32 values and their I32 indices on the last axis. Rank-one and rank-two inputs are admitted. dim may be -1 or the exact last-axis index. Values are sorted descending; ties select lower indices. k is clamped to the last-axis extent, matching OA C++. Errors Returns an error for an unsupported rank, dtype, axis, negative k, a request above the admitted 1024-entry GPU limit, an empty last axis, size overflow, allocation failure, or runtime recording failure.

pub fn top_k_mask(
	indices: &Matrix,
	num_experts: usize
) -> Result<Matrix>

Expand I32 [T, K] expert indices into an exact FP32 [T, E] mask. Errors Returns an error unless the input is nonempty rank-two I32, num_experts is positive, all sizes fit the shader ABI, or runtime recording fails.

pub fn transpose(
	input: &Matrix,
	dim0: i32,
	dim1: i32
) -> Result<Matrix>

Materialize a swap of the last two axes of a rank-two or rank-three Matrix. dim0 and dim1 accept positive or negative axes, but the current donor contract only admits the final two axes in either order. The result owns new storage and preserves the input dtype and batch order. Errors Returns an error unless the input is rank-two or rank-three F32, the axes select the distinct final two dimensions, all geometry fits the shader ABI, allocation succeeds, and runtime recording succeeds.