Machine learning
Models, autograd, losses, optimizers, metrics, and reinforcement learning.
Examples
Functions
Load a native OA .oam checkpoint into existing matching owners. The complete file, section hashes, tensor metadata, module paths, shapes, dtypes, and optimizer payload are validated before any live handle changes. engine remains the sole owner of newly uploaded storage. Errors Returns crate::ErrorKind::CheckpointCorrupt for malformed, truncated, or integrity-invalid .oam data. Returns another error for a valid file that does not match the destination model/optimizer, or for allocation/runtime and filesystem failures.
Save a module and its complete optimizer state in the native OA .oam format. Registration-derived parameter paths become the Weights index. Persistent module buffers become the State index; non-persistent buffers are excluded. This is an explicit host-observation boundary. The completed file is durably written to a sibling temporary path and atomically renamed. Errors Returns an error for ambiguous module ownership, optimizer/model mismatch, unsupported Matrix metadata, readback failure, arithmetic overflow, or a filesystem failure.
Types
1 members
Environment-neutral discrete actor/critic behavior. PPO depends on this contract rather than a concrete network architecture.
Source13 members
Adam optimizer over stable Parameter handles.
Source13 members
Decoupled-weight-decay Adam optimizer over stable Parameter handles.
Source15 members
Deterministic byte-pair encoder compatible with OA’s oa_bpe_v1 files.
Source19 members
Default two-tower categorical actor/critic MLP. Policy and value paths are independent Linear -> ReLU -> Linear -> ReLU -> Linear towers, matching OA C++.
3 members
Dimensions and deterministic initialization seed for the default MLP.
Source12 members
Stateful checkpoint selection and bounded incremental-retention policy.
Source3 members
Filesystem and metric policy for CheckpointManager.
Source0 members
Optimizer with a complete native .oam save-and-restore contract. This capability is sealed because persistence must remain synchronized with OA’s versioned model-file codec. Custom optimizers can implement Optimizer for eager training without claiming wire compatibility.
Source0 members
One bounded continuous action and its differentiable policy statistics.
Source15 members
Checked shape, dtype, and range contract for one environment field. shape excludes the leading environment-batch dimension. An empty shape therefore represents one scalar per environment.
Source11 members
Complete single-agent environment schema shared by native and adapter paths.
Source6 members
One batched environment step result preserving termination and truncation.
Source5 members
Thread-affine reverse-mode recording scope. Constructing a tape selects it for operations on the current thread. GradientTape::backward closes the recording scope and records the corresponding backward operations into the originating engine’s eager batch.
Source1 members
Registers an explicitly included MCP view over an existing training session. The server borrows session through its handlers and must be closed or destroyed first. oa/ml deliberately does not opt applications into the network control surface. Rebuild and preview remain absent until their checkpoint and completion contracts are admitted. Donor: oa::McpTraining.
Source3 members
Gates controlling which MCP tools are registered for a training session. Calling McpTraining::register_tools is the opt-in boundary.
Source15 members
Stateful neural-network component with one registered ownership tree.
Source6 members
Constructor-owned structural registry for one Module. Registration requires &mut self, so a module’s parameter/buffer/child tree is fixed before the module is shared through Rc.
Source13 members
Muon optimizer with Nesterov momentum and rank-two NS5 orthogonalization. Rank-two parameters use the donor Muon matrix pipeline. Other ranks use its fused momentum update; OA never silently delegates them to AdamW.
Source2 members
Owned dotted path and host scalar that participates in module persistence. Scalar state is intended for small algorithmic counters and flags whose value affects the next device operation but does not belong in a Matrix buffer. It is traversed and checkpointed with the same deterministic module tree as parameters and buffers.
Source10 members
Optimizer for externally authored parameter updates. NoOpOptimizer lets callers use crate::ml::ItTraining for cadence, metrics, and callbacks while the step body owns all parameter mutation.
Source6 members
Common behavior exposed by an OA optimizer to training policy. This contract deliberately excludes captured-program and persistence internals. Those capabilities depend on optimizer-specific state and are admitted separately only when their complete implementation exists.
Source7 members
Named trainable FP32 matrix with one live accumulated gradient.
Source7 members
Device-resident replay values, optionally including sampled source indices.
Source10 members
Circular fixed-capacity owner for off-policy transitions. Batched append and deterministic uniform sampling are recorded on the one owning engine. Cursor and size are host control metadata and never require a Matrix readback.
Source3 members
Shape, capacity, and action representation of one replay buffer.
Source8 members
One batch of transitions to append to a ReplayBuffer.
Source11 members
Time-major matrices retained by one RolloutBuffer.
Source11 members
Stateful fixed-capacity owner for one categorical on-policy rollout. Appends, validity reset, and final GAE are recorded on the owning engine; host observation remains the explicit submit-and-wait boundary. Storage is allocated once and overwritten in place across completed collection cycles.
Source3 members
Environment-neutral categorical collector over borrowed owners. Collection records one complete transaction and returns its exact Event without waiting. It owns no environment, model, rollout storage, Engine, or queue.
Source4 members
Fixed categorical collection horizon and sampling policy.
Source3 members
Shape and capacity of one categorical on-policy rollout.
Source10 members
One vector-environment transition whose values remain device-resident.
Source12 members
Enums
Functions
Compute generalized advantage estimates over a fixed rollout. Termination disables value bootstrap. Termination and truncation both stop the reverse trace, while truncation deliberately retains one-step bootstrap. Errors Returns an error unless all inputs share one engine and one nonempty [time, environments] shape, values are F32, boundary masks are U8, both discounts are finite in [0, 1], and GPU recording succeeds.
Record generalized advantage estimation into caller-owned output storage. This is the allocation-free form used by fixed-capacity rollout storage. It has the same semantic identity and numerical contract as gae. Errors Returns the errors from gae, and rejects output matrices that do not match the rollout’s engine, F32 dtype, and shape or alias any input/each other.
Standardize a nonempty FP32 advantage tensor over all elements. The result is (advantage - mean) / sqrt(variance + epsilon). The operation remains device-resident and differentiable, and records one semantic identity over its composite Matrix lowering. Errors Returns an error when advantage is empty or not FP32, epsilon is not positive and finite, a size exceeds the admitted shader ABI, allocation fails, or composite lowering fails.
Types
Functions
Greedily decode one byte from each FP32 logit row. This is an explicit submission, completion, and compact host-read boundary. Errors Returns an error unless logits are rank-one [256] or rank-two [N, 256] FP32, or execution/readback fails.
Greedily decode logits and require exact UTF-8 output. Errors Returns the errors from decode or crate::ErrorKind::DataLoss when the decoded byte stream is not UTF-8.
Upload raw bytes as a rank-one U8 Matrix. Errors Returns an error when allocation or upload fails.
Upload raw bytes as one batched U8 sequence [1, S]. Errors Returns an error when allocation or upload fails.
Upload UTF-8 text without tokenization or normalization. Errors Returns an error when allocation or upload fails.
Sample one byte from each FP32 logit row. Errors Returns an error for an invalid 256-class logit contract, sampling contract, execution, or readback failure.
Constants & Variables
Functions
Clip FP32 rewards to an inclusive finite range. Errors Returns an error for an empty/non-FP32 input, unordered/non-finite bounds, or Matrix/composite lowering failure.
Normalize and symmetrically clip FP32 observations. Computes clamp((observation - mean) / (stddev + epsilon), -clip, clip) through differentiable Core Matrix operations under one Environment semantic identity. Errors Returns an error for empty/non-FP32 inputs, invalid epsilon/clip, incompatible broadcasting or ownership, or composite lowering failure.
Affinely map FP32 actions between two finite ranges. Errors Returns an error for an empty/non-FP32 input, unordered/non-finite bounds, or Matrix/composite lowering failure.
Types
19 members
Stateful native vector-environment behavior. Implementations own semantic environment state and embed exactly one EnvironmentExecution. Public default methods own transaction state, validation boundaries, submission, and rollback; implementations only author reset/step/reset-completed commands and transactional host metadata.
5 members
Private command transaction and exact completion state for one environment. Native environments embed this value and expose it through Environment. It borrows the canonical Engine recorder for each active transaction; it does not own another queue, allocator, scheduler, or execution graph.
Source15 members
Checked shape, dtype, and range contract for one environment field. shape excludes the leading environment-batch dimension. An empty shape therefore represents one scalar per environment.
Source11 members
Complete single-agent environment schema shared by native and adapter paths.
Source6 members
One batched environment step result preserving termination and truncation.
SourceEnums
Functions
Deterministically evaluate a categorical policy and return compact telemetry. The complete horizon is recorded and submitted once, followed by one exact wait and three compact rollout readbacks. Greedy action selection uses the donor temperature-zero route. The environment remains open and idle. Errors Returns an error for zero horizon/lanes, a non-discrete environment, model ownership mismatch, environment/model/policy/rollout recording failure, submission/wait failure, readback failure, or counter overflow.
Types
Functions
Advance one explicit Euler step: state + velocity * delta_time. Errors Returns an error unless state and velocity are matching nonempty same-engine F32 matrices, delta_time is finite, and recording succeeds.
Construct x(t) = clean + t * (noise - clean) and v = noise - clean. A rank-one time vector matching the batch dimension expands across every non-batch axis. Scalars and ordinarily broadcastable time shapes are also accepted. Errors Returns an error unless every input is same-engine F32, clean and noise are matching and nonempty, time broadcasts to them, and recording succeeds.
Compute mean squared error over broadcast-selected valid elements. The mask contributes its values to the denominator. An all-zero mask returns exact zero by clamping the denominator to one. Target and mask are detached; reverse mode differentiates only prediction. Errors Returns an error unless prediction and target are matching nonempty same-engine F32 matrices, mask is same-engine F32 and broadcastable to prediction, and recording succeeds.
Types
Functions
Compute mean binary cross entropy over matching FP32 probability matrices. Predictions use the donor’s [1e-7, 1 - 1e-7] numerical clamp. Targets are treated as detached; reverse mode differentiates only the prediction. Errors Returns an error for empty, mismatched, non-FP32, or cross-engine inputs, or when allocation or runtime recording fails.
Compute mean cross-entropy over rank-two FP32 logits and U32 or I32 class targets. logits has shape [N, C], targets has shape [N], and the returned matrix is an FP32 scalar. I32 targets must be non-negative; any negative or out-of-range target produces NaN without an out-of-bounds access. Errors Returns an error when ranks, shapes, dtypes, ownership, allocation, or runtime recording are invalid.
Compute mean cross-entropy over the rows selected by a floating FP32 mask. logits has shape [N, C]; targets and mask have shape [N]. A zero mask value excludes its row and produces an exact-zero logits adjoint. valid_count is the caller-provided normalization denominator and must be in 1..=N; it is not read back or recomputed from device storage. Reverse mode differentiates only logits. Errors Returns an error when ranks, shapes, dtypes, ownership, valid_count, allocation, or runtime recording are invalid.
Compose the donor DQN selected-action and detached Bellman-target loss. Termination suppresses next-state bootstrap. Truncation deliberately keeps bootstrap because it ends only the current trace, not the underlying task. Reverse mode therefore connects loss and selected_q only to q. Errors Returns an error unless Q and next-Q are same-engine nonempty F32 [B,A] matrices with A > 1, action is I32 [B], reward is F32 [B], boundary masks are U8 [B], all fields share one engine, extents fit the shader ABI, and discount is finite in [0, 1].
Compute mean absolute error for matching FP32 matrices. The target is treated as detached; reverse mode differentiates only the prediction. Equal prediction and target values use the zero subgradient. Errors Returns an error for empty, mismatched, non-FP32, or cross-engine inputs, or when allocation or runtime recording fails.
Compute mean squared error for matching FP32 matrices. The target is treated as detached; reverse mode differentiates only the prediction. Errors Returns an error for empty, mismatched, non-FP32, or cross-engine inputs, or when allocation or runtime recording fails.
Compose PPO policy, value, entropy, and total losses. All four outputs are FP32 scalars under one semantic PPO operation. Old log probability, advantage, and target return remain detached by their child operation contracts. Errors Returns an error unless all rollout fields are matching nonempty same-engine FP32 matrices and every coefficient is valid.
Compute the mean negative clipped PPO surrogate objective. Errors Returns an error unless all inputs are matching nonempty same-engine FP32 matrices and clip_epsilon is finite in (0, 1).
Compute the explicit clipped-policy adjoint for new log probability. Errors Returns the same validation and runtime errors as ppo_clipped_policy.
Compute the donor entropy-regularized SAC actor objective. The scalar result is mean(alpha * log_probability - min(q1, q2)) and is differentiable through all three inputs. Errors Returns an error unless inputs are matching nonempty same-engine F32 vectors, entropy coefficient is finite and non-negative, and lowering succeeds.
Compose the donor twin-critic SAC objective. The target is reward + discount * (min(next_q1, next_q2) - alpha * next_log_probability) outside terminal states. Truncation retains that bootstrap. All next-state target fields are detached from reverse mode. Errors Returns an error unless all numerical fields are matching nonempty same-engine F32 vectors, both boundary masks are matching U8 vectors, discount is finite in [0, 1], entropy coefficient is non-negative and finite, the batch fits the shader ABI, and lowering succeeds.
Compute mean unit-beta Smooth L1 loss for matching FP32 matrices. The target is treated as detached; reverse mode differentiates only the prediction. Errors Returns an error for empty, mismatched, non-FP32, or cross-engine inputs, or when allocation or runtime recording fails.
Types
Functions
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.
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.
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.
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.
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.
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.
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.
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.
Compute complete explicit channel-normalization adjoints. Errors Returns an error when the forward contract or output gradient is invalid, or runtime recording fails.
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.
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.
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.
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.
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.
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.
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.
Apply the exponential linear unit elementwise. Errors Returns an error unless input is F32 or runtime recording fails.
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.
Apply the tanh-approximate Gaussian error linear unit elementwise. Errors Returns an error unless input is F32 or runtime recording fails.
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.
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.
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.
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.
Apply Leaky ReLU elementwise with the requested negative slope. Errors Returns an error unless input is F32 or runtime recording fails.
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.
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.
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.
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.
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.
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.
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.
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.
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.
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.
Apply Mish elementwise. Errors Returns an error unless input is F32 or runtime recording fails.
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.
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.
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.
Apply the rectified linear unit elementwise. Errors Returns an error unless input is F32 or runtime recording fails.
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.
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.
Compute all gated RMS normalization adjoints explicitly. Errors Returns an error for an incompatible output gradient or forward contract, or when runtime recording fails.
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.
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.
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.
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.
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.
Apply the logistic sigmoid elementwise. Errors Returns an error unless input is F32 or runtime recording fails.
Apply the sigmoid linear unit elementwise. Errors Returns an error unless input is F32 or runtime recording fails.
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.
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.
Apply the overflow-safe Softplus formulation elementwise. Errors Returns an error unless input is F32 or runtime recording fails.
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.
Apply SiLU(gate) * up elementwise. Errors Returns an error unless both inputs are equal-shape F32 on the same engine, or runtime recording fails.
Apply hyperbolic tangent elementwise. Errors Returns an error unless input is F32 or runtime recording fails.
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.
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.
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.
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.
Types
0 members
The fifteen explicit adjoints of mamba3_mimo.
Source3 members
Geometry and numerical bounds for the fused Mamba-3 projection preprocess.
Source1 members
Eight values produced by the fused Mamba-3 projection preprocess.
Source0 members
Complete adjoints of broadcast-affine gated RMS normalization.
Source0 members
The eleven explicit adjoints of mamba3_siso.
Source3 members
Explicit tensor geometry and optional paths for a selective-state scan. A zero num_groups means one Q/K group per head, matching OA C++.
SourceEnums
Types
1 members
Registers an explicitly included MCP view over an existing training session. The server borrows session through its handlers and must be closed or destroyed first. oa/ml deliberately does not opt applications into the network control surface. Rebuild and preview remain absent until their checkpoint and completion contracts are admitted. Donor: oa::McpTraining.
Source3 members
Gates controlling which MCP tools are registered for a training session. Calling McpTraining::register_tools is the opt-in boundary.
SourceFunctions
Compute last-axis categorical accuracy as a fraction in [0, 1]. Argmax and counting execute on the owning device. Only the resulting U32 scalar crosses the host boundary; this function then waits for that scalar. Equal maximum logits select the first class, matching the OA C++ contract. Errors Returns an error when matrix::categorical_accuracy_count rejects the logits or labels, or submission, completion, or scalar readback fails.
Read one scalar FP32 loss value. This is an explicit host-observation boundary: pending work is submitted and the producing event is awaited before the value is returned. Errors Returns an error unless loss contains exactly one F32 value, or submission, completion, or readback fails.
Types
20 members
21 members
27 members
20 members
20 members
27 members
27 members
25 members
26 members
19 members
22 members
Trainable FP32 lookup table indexed by an arbitrary-shape U8, U32, or I32 matrix.
Construct a deterministically initialized embedding table. Errors Returns an error when either dimension is zero, shape arithmetic overflows, or allocation/upload fails.
Construct a deterministically initialized embedding table. Errors Returns an error when either dimension is zero, shape arithmetic overflows, or allocation/upload fails.
Construct a deterministically initialized embedding table. Errors Returns an error when either dimension is zero, shape arithmetic overflows, or allocation/upload fails.
21 members
OA’s embedding-plus-Mamba core with one canonical Rust execution path. The C++ donor currently carries renamed Empyrealm shader copies whose SPIR-V is identical to the Mamba-3 providers. Rust preserves the public module and parameter topology while reusing Mamba3 until the algorithms actually diverge.
20 members
Pre-normalized SwiGLU feed-forward network with a residual connection. The module preserves the OA donor structure: RMSNorm -> (gate, up) -> SwiGLU -> down -> residual add.
20 members
Flatten an inclusive range of Matrix dimensions into one dimension. Negative dimensions are resolved from the end at call time. The donor default (1, -1) preserves a leading batch dimension.
23 members
Trainable flow/diffusion denoiser with optional condition features and CFG.
Construct a deterministic dense or MoE denoiser. Errors Returns an error for invalid model/time/dropout geometry or failed child, parameter, and buffer construction.
Construct a deterministic dense or MoE denoiser. Errors Returns an error for invalid model/time/dropout geometry or failed child, parameter, and buffer construction.
Construct a deterministic dense or MoE denoiser. Errors Returns an error for invalid model/time/dropout geometry or failed child, parameter, and buffer construction.
4 members
Configuration for a modality-independent flow/diffusion denoiser.
Source20 members
24 members
Reusable bidirectional Transformer backbone for flow/diffusion denoisers.
Construct a deterministic dense or dropless-MoE flow Transformer. Errors Returns an error for invalid geometry, expert routing, epsilon, or failed child construction and registration.
Construct a deterministic dense or dropless-MoE flow Transformer. Errors Returns an error for invalid geometry, expert routing, epsilon, or failed child construction and registration.
Construct a deterministic dense or dropless-MoE flow Transformer. Errors Returns an error for invalid geometry, expert routing, epsilon, or failed child construction and registration.
4 members
Configuration for a bidirectional flow/diffusion Transformer backbone.
Source18 members
25 members
Stacked batch-first gated recurrent unit using whole-sequence GPU scans. Input uses [batch, sequence, input]; output uses [batch, sequence, hidden]. Every layer starts from a zero hidden state.
24 members
18 members
25 members
26 members
19 members
24 members
4 members
Construction contract for the donor-backed Mamba-3 block.
Source1 members
Explicit recurrent cache for one Mamba3 module and batch size. The cache owns the selective state, cumulative rotary angle, previous key, and previous value buffers. It is deliberately separate from the module so independent decoding streams never share hidden mutable state.
Source21 members
43 members
Sparse top-k mixture-of-experts feed-forward module. The implementation preserves OA’s expert-major route plan and stacked parameter layout. Routing decisions remain detached while gradients flow through selected gate magnitudes, expert projections, RMSNorm, and the residual input.
8 members
Reduced host telemetry for the latest MoE routing decision.
Source31 members
Multi-head self-attention over packed [B*S, D] values.
Construct four biased projections with causal visibility and automatic backend policy. Errors Returns an error unless dimensions are positive and model_width is divisible by num_heads, or child construction/registration fails.
Construct four biased projections with causal visibility and automatic backend policy. Errors Returns an error unless dimensions are positive and model_width is divisible by num_heads, or child construction/registration fails.
18 members
23 members
22 members
26 members
Stacked batch-first Elman recurrent network using whole-sequence GPU scans. Input uses [batch, sequence, input]; output uses [batch, sequence, hidden]. Every layer starts from a zero hidden state.
24 members
20 members
22 members
18 members
19 members
20 members
26 members
Ready-to-train causal Transformer language model. This is the idiomatic Rust spelling of donor oa::NnTransformer: token and position embeddings, registered Transformer blocks, final normalization, and a vocabulary projection.
Construct a deterministic causal language model. Errors Returns an error unless all dimensions are positive, the model width is divisible by the head count, positions fit U32, and every child can be constructed and registered.
Construct a deterministic causal language model. Errors Returns an error unless all dimensions are positive, the model width is divisible by the head count, positions fit U32, and every child can be constructed and registered.
Construct a deterministic causal language model. Errors Returns an error unless all dimensions are positive, the model width is divisible by the head count, positions fit U32, and every child can be constructed and registered.
Construct a deterministic causal language model. Errors Returns an error unless all dimensions are positive, the model width is divisible by the head count, positions fit U32, and every child can be constructed and registered.
34 members
One pre-normalized causal Transformer block.
Construct the canonical pre-norm attention and GELU feed-forward block. Errors Returns an error for invalid dimensions or failed child construction.
Construct the canonical pre-norm attention and GELU feed-forward block. Errors Returns an error for invalid dimensions or failed child construction.
20 members
24 members
4 members
Configuration for one EMA-trained vector-quantization codebook.
SourceEnums
Execution policy for scaled dot-product attention. Auto selects only a route proven by device-specific performance evidence. It currently preserves the compositional AttentionBackend::Standard provider on every device.
Auto · Standard · Flash
Functions
Clip the combined L2 norm of nonempty FP32 gradients in place. The operation records one GPU reduction followed by one GPU scaling pass. It does not read gradients back to the CPU or wait for completion. Empty matrices are ignored; at most sixteen nonempty gradients are admitted. Errors Returns an error when max_norm is negative or non-finite, more than sixteen nonempty gradients are supplied, gradients do not share one engine, a gradient is not FP32, storage is duplicated, a size exceeds the shader ABI, or allocation and runtime recording fail.
Types
13 members
Adam optimizer over stable Parameter handles.
Source13 members
Decoupled-weight-decay Adam optimizer over stable Parameter handles.
Source0 members
Optimizer with a complete native .oam save-and-restore contract. This capability is sealed because persistence must remain synchronized with OA’s versioned model-file codec. Custom optimizers can implement Optimizer for eager training without claiming wire compatibility.
Source13 members
Muon optimizer with Nesterov momentum and rank-two NS5 orthogonalization. Rank-two parameters use the donor Muon matrix pipeline. Other ranks use its fused momentum update; OA never silently delegates them to AdamW.
Source10 members
Optimizer for externally authored parameter updates. NoOpOptimizer lets callers use crate::ml::ItTraining for cadence, metrics, and callbacks while the step body owns all parameter mutation.
Source6 members
Common behavior exposed by an OA optimizer to training policy. This contract deliberately excludes captured-program and persistence internals. Those capabilities depend on optimizer-specific state and are admitted separately only when their complete implementation exists.
Source12 members
Functions
Differentiably evaluate stored categorical actions. Gradients flow through logits, log_probability, and entropy; action indices remain discrete. action and value are zero-copy semantic output aliases, matching OA C++. Errors Returns an error when logits are not FP32 [E,A] with E > 0 and A > 1, actions are not I32 [E], values are not FP32 [E], inputs do not share one engine, or composite lowering fails.
Differentiably evaluate a stored pre-tanh normal action. Log standard deviation is clamped to [-20, 2]. Actions are mapped from [-1, 1] into [minimum, maximum]; log probability includes the donor tanh Jacobian correction, while entropy is the unsquashed diagonal-normal value. raw_action and value are zero-copy semantic output aliases. Errors Returns an error when mean, log-standard-deviation, and raw-action are not matching nonempty FP32 [E,A] matrices, value is not FP32 [E], inputs do not share one engine, bounds/epsilon are invalid, a size exceeds the shader ABI, or composite lowering fails.
Sample and evaluate one categorical action per environment. A nonzero seed is deterministic. Seed zero follows the process-local OA RNG convention. The action is sampled from the same logits used for returned log probabilities and entropy. Errors Returns an error when logits are not FP32 [E,A] with E > 0 and A > 1, values are not FP32 [E] on the same engine, or composite lowering fails.
Sample a reparameterized tanh-normal policy. The returned raw action is mean + exp(clamped_log_stddev) * noise, so gradients flow through both distribution parameters while Philox noise stays detached. A nonzero seed is deterministic; seed zero follows OA’s process-local RNG convention. Errors Returns an error when mean/log-standard-deviation are not matching nonempty FP32 [E,A] matrices, value is not FP32 [E] on the same engine, bounds/epsilon are invalid, a size exceeds the shader ABI, or composite lowering fails.
Types
Types
7 members
Device-resident replay values, optionally including sampled source indices.
Source10 members
Circular fixed-capacity owner for off-policy transitions. Batched append and deterministic uniform sampling are recorded on the one owning engine. Cursor and size are host control metadata and never require a Matrix readback.
Source3 members
Shape, capacity, and action representation of one replay buffer.
Source8 members
One batch of transitions to append to a ReplayBuffer.
SourceTypes
11 members
Time-major matrices retained by one RolloutBuffer.
Source11 members
Stateful fixed-capacity owner for one categorical on-policy rollout. Appends, validity reset, and final GAE are recorded on the owning engine; host observation remains the explicit submit-and-wait boundary. Storage is allocated once and overwritten in place across completed collection cycles.
Source3 members
Shape and capacity of one categorical on-policy rollout.
Source10 members
One vector-environment transition whose values remain device-resident.
SourceTypes
Types
0 members
OA C++ compatibility spelling for Checkpoint.
Source0 members
OA C++ compatibility spelling for CsvLogger.
Source0 members
OA C++ compatibility spelling for EarlyStopping.
Source0 members
OA C++ compatibility spelling for LearningRateScheduler.
Source0 members
OA C++ compatibility spelling for PhaseSchedule.
Source0 members
OA C++ compatibility spelling for ProgressBar.
Source0 members
OA C++ compatibility spelling for TrainingSummary.
Source0 members
OA C++ compatibility spelling for Validation.
Source11 members
Epoch, interval, best-model, and restore-best checkpoint policy.
Source7 members
Buffered one-row-per-completed-step CSV logger. File creation, writes, and close-time flushing happen only at explicit callback boundaries. Abandoning the callback discards any unflushed host buffer; Drop does not hide an I/O failure.
Source10 members
Stateful DQN training coordinator over caller-owned components. The online module, target module, optimizer, replay storage, and engine remain caller-owned and must outlive this coordinator. Dropping the trainer does not submit, wait, synchronize targets, or fire terminal callbacks; use DqnTrainer::finish for the explicit terminal boundary.
Source4 members
Fixed update budget, replay batch, target cadence, and Bellman policy.
Source12 members
Cooperative epoch-boundary early stopping.
Source4 members
Aggregate device time for completed, timestamped training steps.
Source14 members
Coordinator for alternating fixed-capacity collection and optimizer phases.
Source26 members
Exact eager or captured training-step lifecycle. The loop borrows the engine, optimizer, metrics, and callbacks. Call ItTraining::begin_step before authoring each body, then finish it with ItTraining::complete_step or ItTraining::complete_program_step. ItTraining::finish is the explicit fallible terminal boundary.
Construct a training lifecycle around one engine and optimizer. Errors Returns an error when source-unit accounting is non-finite or overflows, or the variable epoch schedule overflows.
Construct a training lifecycle around one engine and optimizer. Errors Returns an error when source-unit accounting is non-finite or overflows, or the variable epoch schedule overflows.
Construct a training lifecycle around one engine and optimizer. Errors Returns an error when source-unit accounting is non-finite or overflows, or the variable epoch schedule overflows.
Construct a training lifecycle around one engine and optimizer. Errors Returns an error when source-unit accounting is non-finite or overflows, or the variable epoch schedule overflows.
Construct a training lifecycle around one engine and optimizer. Errors Returns an error when source-unit accounting is non-finite or overflows, or the variable epoch schedule overflows.
Construct a training lifecycle around one engine and optimizer. Errors Returns an error when source-unit accounting is non-finite or overflows, or the variable epoch schedule overflows.
3 members
Configuration for one explicit training loop.
Source6 members
Applies a learning-rate schedule after every completed optimizer step.
Source13 members
Consecutive phase policy over one super::super::ItTraining lifecycle. The phase schedule does not own another iterator. Its epoch/step totals must exactly match the iterator configuration, and its optional transition hook receives restricted mutable access to the iterator’s existing optimizer.
Source17 members
Complete caller-driven categorical PPO lifecycle. The Engine, model, and optimizer remain caller-owned. The trainer owns only its fixed-capacity rollout and lifecycle state. Drop performs no submission, waiting, or terminal callbacks; call PpoTrainer::finish explicitly.
4 members
Complete categorical PPO collection and update policy.
Source4 members
Scalar results from the last completed PPO update epoch.
Source9 members
Throttled Keras/tqdm-style wall-throughput progress display.
Source10 members
Stateful fixed-temperature SAC coordinator over caller-owned components.
Source4 members
Fixed SAC update, shape, action-range, target, and loss policy.
Source5 members
Restricted mutable services available to one callback invocation.
Source3 members
One command submitted by a session client.
Source3 members
Durable result of one processed command.
Source7 members
Immutable evidence for one training-program compilation stage.
Source0 members
Compatibility alias for the pre-ItTraining Rust prototype spelling.
Source0 members
Compatibility alias for the pre-ItTrainingConfig Rust prototype spelling.
Source17 members
Fixed-shape forward, backward, and AdamW program captured for Vulkan replay. The first capture records semantic work without executing it. Replays use one cached command buffer and graph-resident optimizer state; no forward, autograd, or optimizer graph is rebuilt on the host.
20 members
Cloneable, thread-safe command and observation handle for one training loop. Attach it to exactly one super::ItTraining instance. Clones may enqueue commands and read bounded evidence from UI, viewer, or service threads.
3 members
Bounded live-control snapshot published by an attached iterator.
Source28 members
9 members
Final loss, wall-time, device-time, throughput, and run summary.
Source9 members
Runs application-owned inference validation at completed lifecycle boundaries. Epoch-based training evaluates once at every epoch end. Step-only training evaluates at each non-zero configured interval and once at train end when the final step was not already evaluated. Evaluation wall time is excluded from training throughput and the callback never submits work on its own.
Source4 members
Cloneable observation handle for the latest completed validation loss. Clones share one thread-affine value so checkpoint, summary, and early-stop callbacks can observe the result without owning or rerunning validation.
SourceEnums
Kind of command applied at a training safe point.
Pause · Resume · Stop · Checkpoint · Evaluate · SetParameter · RequestRecapture · RequestRebuild
Ordered compilation stage for a captured training program.
SemanticValidation · ReplaySafety · Decomposition · Fusion · Placement · Precision · KernelSelection · LoweringValidation · MemoryPlanning · SynchronizationPlanning · CommandRecording
Outcome of one training-program compilation stage.
NotRun · Inherited · Analyzed · Applied · Failed
Live state of an attached training lifecycle.
Running · Paused · Stopping · Completed · Failed
Typed value carried by a command or exposed parameter.
Empty · Boolean(bool) · Integer(i64) · Float(f64) · String(String)