nuvec_nn/
losses.rs

1//! Loss functions, as in MLX's `nn.losses`.
2
3use nuvec::{Array, DType, fast, ops};
4
5/// How per-example losses are combined.
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum Reduction {
8    /// One loss per example.
9    None,
10    /// The mean over examples.
11    Mean,
12    /// The sum over examples.
13    Sum,
14}
15
16/// Cross entropy of `logits` (classes on the last axis) against integer class indices
17/// `targets` (the logits' shape without the last axis), as MLX's `nn.losses.cross_entropy`:
18/// `logsumexp(z) − z[target]` with `z = logits − max(logits)` (the maximum without
19/// gradient); on the GPU, the fused [`nuvec::fast::cross_entropy`] in the logits' dtype.
20///
21/// Mismatched shapes or non-integer targets give a failed array (with the error of
22/// [`nuvec::fast::cross_entropy`]).
23#[track_caller]
24pub fn cross_entropy(logits: &Array, targets: &Array, reduction: Reduction) -> Array {
25    let valid = logits.ndim() > 0
26        && targets.shape() == &logits.shape()[..logits.ndim() - 1]
27        && is_index_dtype(targets.dtype());
28    let loss = if !valid || nuvec::default_device().kind == nuvec::DeviceKind::Gpu {
29        // One fused kernel (and one for the gradient), as MLX on the GPU; it also reports
30        // invalid arguments as a failed array.
31        fast::cross_entropy(logits, targets).astype(logits.dtype())
32    } else {
33        let classes = logits.shape()[logits.ndim() - 1];
34        let z = logits.subtract(logits.max_keepdims(-1).stop_gradient());
35        let positions = ops::arange(0, classes, 1).astype(targets.dtype());
36        let hit = positions.equal(targets.expand_dims(-1));
37        let score = hit.select(&z, 0.0).sum(-1);
38        z.logsumexp(-1).subtract(&score)
39    };
40    match reduction {
41        Reduction::None => loss,
42        Reduction::Mean => loss.mean(..),
43        Reduction::Sum => loss.sum(..),
44    }
45}
46
47/// Class indices are integers (not one-hot floats).
48fn is_index_dtype(d: DType) -> bool {
49    !matches!(
50        d,
51        DType::F16 | DType::BF16 | DType::F32 | DType::F64 | DType::Complex64
52    )
53}
nuvec_nn/losses.rs

nuvec_nn/
losses.rs
