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
#[allow(clippy::wildcard_imports)]
use super::*;
impl NLLLoss {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn forward(&self, log_probs: &Tensor, targets: &Tensor) -> Tensor {
assert_eq!(log_probs.ndim(), 2);
assert_eq!(targets.ndim(), 1);
let batch_size = log_probs.shape()[0];
let num_classes = log_probs.shape()[1];
let mut losses = Vec::with_capacity(batch_size);
for b in 0..batch_size {
let target = targets.data()[b] as usize;
losses.push(-log_probs.data()[b * num_classes + target]);
}
let loss = Tensor::new(&losses, &[batch_size]);
match self.reduction {
Reduction::None => loss,
Reduction::Mean => loss.mean(),
Reduction::Sum => loss.sum(),
}
}
}
// ============================================================================
// Helper Functions
// ============================================================================
/// Element-wise absolute value.
///
/// PMAT-896: delegates to the autograd-aware `Tensor::abs` so the
/// computation graph stays connected. The previous `Tensor::from_vec`
/// implementation severed autograd, which silently zeroed every
/// `L1Loss` gradient (`L1Loss = mean(|pred - target|)`).
pub(super) fn abs(x: &Tensor) -> Tensor {
x.abs()
}
/// Softmax computation for gradient tracking.
///
/// ONE PATH: Delegates to `nn::functional::softmax` (UCBD §4).
pub(super) fn softmax_2d(x: &Tensor) -> Tensor {
crate::nn::functional::softmax(x, -1)
}
/// Log-softmax for numerical stability.
///
/// ONE PATH: Delegates to `nn::functional::log_softmax` (UCBD §4).
pub(super) fn log_softmax(x: &Tensor) -> Tensor {
crate::nn::functional::log_softmax(x, -1)
}