nuvec-nn 0.2.0

Neural-network modules, layers and optimizers for nuvec
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64

//! Loss functions, as in MLX's `nn.losses`.

use nuvec::{
    Array, DType, Error, Result, arange, expand_dims, logsumexp, max, mean, stop_gradient,
    subtract, sum, where_,
};

/// How per-example losses are combined.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Reduction {
    /// One loss per example.
    None,
    /// The mean over examples.
    Mean,
    /// The sum over examples.
    Sum,
}

/// Cross entropy of `logits` (classes on the last axis) against integer class indices
/// `targets` (the logits' shape without the last axis), as MLX's `nn.losses.cross_entropy`:
/// `logsumexp(z) − z[target]` with `z = logits − max(logits)` (the maximum without
/// gradient); on the GPU, the fused [`nuvec::fast::cross_entropy`] in the logits' dtype.
pub fn cross_entropy(logits: &Array, targets: &Array, reduction: Reduction) -> Result<Array> {
    let classes = *logits
        .shape()
        .last()
        .ok_or_else(|| Error::InvalidArgument("cross_entropy: logits need a class axis".into()))?;
    if targets.shape() != &logits.shape()[..logits.ndim() - 1] {
        return Err(Error::InvalidArgument(format!(
            "cross_entropy: targets shape {:?} does not match logits shape {:?}",
            targets.shape(),
            logits.shape()
        )));
    }
    if !is_index_dtype(targets.dtype()) {
        return Err(Error::InvalidArgument(format!(
            "cross_entropy: targets must be integer class indices, got {}",
            targets.dtype()
        )));
    }
    let loss = if nuvec::default_device().kind == nuvec::DeviceKind::Gpu {
        // One fused kernel (and one for the gradient), as MLX on the GPU.
        nuvec::fast::cross_entropy(logits, targets)?.astype(logits.dtype())?
    } else {
        let z = subtract(logits, stop_gradient(&max(logits, &[-1], true)?)?)?;
        let positions = arange(0.0, classes as f64, 1.0, targets.dtype())?;
        let hit = nuvec::equal(&positions, expand_dims(targets, -1)?)?;
        let score = sum(&where_(&hit, &z, 0.0)?, &[-1], false)?;
        subtract(logsumexp(&z, &[-1], false)?, &score)?
    };
    match reduction {
        Reduction::None => Ok(loss),
        Reduction::Mean => mean(&loss, &[], false),
        Reduction::Sum => sum(&loss, &[], false),
    }
}

/// Class indices are integers (not one-hot floats).
fn is_index_dtype(d: DType) -> bool {
    !matches!(
        d,
        DType::F16 | DType::BF16 | DType::F32 | DType::F64 | DType::Complex64
    )
}