nuvec_nn/
losses.rs

1//! Loss functions, as in MLX's `nn.losses`.
2
3use nuvec::{
4    Array, DType, Error, Result, arange, expand_dims, logsumexp, max, mean, stop_gradient,
5    subtract, sum, where_,
6};
7
8/// How per-example losses are combined.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum Reduction {
11    /// One loss per example.
12    None,
13    /// The mean over examples.
14    Mean,
15    /// The sum over examples.
16    Sum,
17}
18
19/// Cross entropy of `logits` (classes on the last axis) against integer class indices
20/// `targets` (the logits' shape without the last axis), as MLX's `nn.losses.cross_entropy`:
21/// `logsumexp(z) − z[target]` with `z = logits − max(logits)` (the maximum without
22/// gradient); on the GPU, the fused [`nuvec::fast::cross_entropy`] in the logits' dtype.
23pub fn cross_entropy(logits: &Array, targets: &Array, reduction: Reduction) -> Result<Array> {
24    let classes = *logits
25        .shape()
26        .last()
27        .ok_or_else(|| Error::InvalidArgument("cross_entropy: logits need a class axis".into()))?;
28    if targets.shape() != &logits.shape()[..logits.ndim() - 1] {
29        return Err(Error::InvalidArgument(format!(
30            "cross_entropy: targets shape {:?} does not match logits shape {:?}",
31            targets.shape(),
32            logits.shape()
33        )));
34    }
35    if !is_index_dtype(targets.dtype()) {
36        return Err(Error::InvalidArgument(format!(
37            "cross_entropy: targets must be integer class indices, got {}",
38            targets.dtype()
39        )));
40    }
41    let loss = if nuvec::default_device().kind == nuvec::DeviceKind::Gpu {
42        // One fused kernel (and one for the gradient), as MLX on the GPU.
43        nuvec::fast::cross_entropy(logits, targets)?.astype(logits.dtype())?
44    } else {
45        let z = subtract(logits, stop_gradient(&max(logits, &[-1], true)?)?)?;
46        let positions = arange(0.0, classes as f64, 1.0, targets.dtype())?;
47        let hit = nuvec::equal(&positions, expand_dims(targets, -1)?)?;
48        let score = sum(&where_(&hit, &z, 0.0)?, &[-1], false)?;
49        subtract(logsumexp(&z, &[-1], false)?, &score)?
50    };
51    match reduction {
52        Reduction::None => Ok(loss),
53        Reduction::Mean => mean(&loss, &[], false),
54        Reduction::Sum => sum(&loss, &[], false),
55    }
56}
57
58/// Class indices are integers (not one-hot floats).
59fn is_index_dtype(d: DType) -> bool {
60    !matches!(
61        d,
62        DType::F16 | DType::BF16 | DType::F32 | DType::F64 | DType::Complex64
63    )
64}
nuvec_nn/losses.rs

nuvec_nn/
losses.rs
