oa::ml::loss

Module published by the oa::ml module.

MdRust public API

Public Types

Public Functions

pub fn bce(prediction: &Matrix, target: &Matrix) -> Result<Matrix>

pub fn cross_entropy(logits: &Matrix, targets: &Matrix) -> Result<Matrix>

pub fn cross_entropy_masked( logits: &Matrix, targets: &Matrix, mask: &Matrix, valid_count: usize, ) -> Result<Matrix>

pub fn dqn( q: &Matrix, action: &Matrix, reward: &Matrix, next_q: &Matrix, terminated: &Matrix, truncated: &Matrix, config: DqnLossConfig, ) -> Result<DqnLossResult>

pub fn l1(prediction: &Matrix, target: &Matrix) -> Result<Matrix>

pub fn mse(prediction: &Matrix, target: &Matrix) -> Result<Matrix>

pub fn ppo( new_log_probability: &Matrix, old_log_probability: &Matrix, advantage: &Matrix, value: &Matrix, target_return: &Matrix, entropy: &Matrix, config: PpoLossConfig, ) -> Result<PpoLossResult>

pub fn ppo_clipped_policy( new_log_probability: &Matrix, old_log_probability: &Matrix, advantage: &Matrix, clip_epsilon: f32, ) -> Result<Matrix>

pub fn ppo_clipped_policy_backward( new_log_probability: &Matrix, old_log_probability: &Matrix, advantage: &Matrix, clip_epsilon: f32, ) -> Result<Matrix>

pub fn sac_actor( q1: &Matrix, q2: &Matrix, log_probability: &Matrix, entropy_coefficient: f32, ) -> Result<Matrix>

pub fn sac_critic( q1: &Matrix, q2: &Matrix, reward: &Matrix, next_q1: &Matrix, next_q2: &Matrix, next_log_probability: &Matrix, terminated: &Matrix, truncated: &Matrix, config: SacLossConfig, ) -> Result<SacCriticLossResult>

pub fn smooth_l1(prediction: &Matrix, target: &Matrix) -> Result<Matrix>

Function Documentation

pub fn bce(
	prediction: &Matrix,
	target: &Matrix
) -> Result<Matrix>

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.

pub fn cross_entropy(
	logits: &Matrix,
	targets: &Matrix
) -> Result<Matrix>

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.

pub fn cross_entropy_masked(
	logits: &Matrix,
	targets: &Matrix,
	mask: &Matrix,
	valid_count: usize
) -> Result<Matrix>

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.

pub fn dqn(
	q: &Matrix,
	action: &Matrix,
	reward: &Matrix,
	next_q: &Matrix,
	terminated: &Matrix,
	truncated: &Matrix,
	config: DqnLossConfig
) -> Result<DqnLossResult>

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

pub fn l1(
	prediction: &Matrix,
	target: &Matrix
) -> Result<Matrix>

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.

pub fn mse(
	prediction: &Matrix,
	target: &Matrix
) -> Result<Matrix>

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.

pub fn ppo(
	new_log_probability: &Matrix,
	old_log_probability: &Matrix,
	advantage: &Matrix,
	value: &Matrix,
	target_return: &Matrix,
	entropy: &Matrix,
	config: PpoLossConfig
) -> Result<PpoLossResult>

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.

pub fn ppo_clipped_policy(
	new_log_probability: &Matrix,
	old_log_probability: &Matrix,
	advantage: &Matrix,
	clip_epsilon: f32
) -> Result<Matrix>

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

pub fn ppo_clipped_policy_backward(
	new_log_probability: &Matrix,
	old_log_probability: &Matrix,
	advantage: &Matrix,
	clip_epsilon: f32
) -> Result<Matrix>

Compute the explicit clipped-policy adjoint for new log probability. Errors Returns the same validation and runtime errors as ppo_clipped_policy.

pub fn sac_actor(
	q1: &Matrix,
	q2: &Matrix,
	log_probability: &Matrix,
	entropy_coefficient: f32
) -> Result<Matrix>

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.

pub fn sac_critic(
	q1: &Matrix,
	q2: &Matrix,
	reward: &Matrix,
	next_q1: &Matrix,
	next_q2: &Matrix,
	next_log_probability: &Matrix,
	terminated: &Matrix,
	truncated: &Matrix,
	config: SacLossConfig
) -> Result<SacCriticLossResult>

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.

pub fn smooth_l1(
	prediction: &Matrix,
	target: &Matrix
) -> Result<Matrix>

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.