nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

nll_loss to mixin (#15918)

Browse Source
This commit is contained in:
chenyu
2026-04-24 15:50:31 -04:00
committed by GitHub
parent 2f9fdb4a37
commit b501ba3e42
4 changed files with 43 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_tensor_uop_mixin.py
+16
View File
@@ -163,6 +163,22 @@ class TestTensorUOpLoss(unittest.TestCase):
t, Y = _t(2, 3).float(), Tensor([1, 2], dtype=dtypes.int32)
self.assertIs(_strip_unique(t.sparse_categorical_crossentropy(Y, ignore_index=0).uop),
_strip_unique(t.uop.sparse_categorical_crossentropy(Y.uop, ignore_index=0)))
def test_nll_loss(self):
t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
self.assertIs(_strip_unique(t.nll_loss(Y).uop), _strip_unique(t.uop.nll_loss(Y.uop)))
def test_nll_loss_weight(self):
t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float()
self.assertIs(_strip_unique(t.nll_loss(Y, weight=w).uop), _strip_unique(t.uop.nll_loss(Y.uop, weight=w.uop)))
def test_nll_loss_ignore_index(self):
t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
self.assertIs(_strip_unique(t.nll_loss(Y, ignore_index=1).uop), _strip_unique(t.uop.nll_loss(Y.uop, ignore_index=1)))
def test_nll_loss_none_reduction(self):
t, Y = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32)
self.assertIs(_strip_unique(t.nll_loss(Y, reduction="none").uop), _strip_unique(t.uop.nll_loss(Y.uop, reduction="none")))
def test_nll_loss_weight_ignore_index(self):
t, Y, w = _t(2, 3).float().log_softmax(), Tensor([1, 2], dtype=dtypes.int32), _t(3).float()
self.assertIs(_strip_unique(t.nll_loss(Y, weight=w, ignore_index=1).uop),
_strip_unique(t.uop.nll_loss(Y.uop, weight=w.uop, ignore_index=1)))
class TestTensorUOpScatter(unittest.TestCase):
def test_scatter(self):
tinygrad/mixin/__init__.py
+24
View File
@@ -1318,6 +1318,30 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
Y = (1 - label_smoothing)*Y + label_smoothing / int(Y.shape[classes_dim])
return -self.log_softmax(classes_dim).mul(Y).sum(classes_dim)._do_reduction(reduction)
def nll_loss(self, Y:Self, weight:Self|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean") -> Self:
"""
Computes the negative log likelihood loss between log-probabilities and target labels.
NOTE: `self` is log-probabilities and `Y` is the Y labels or class probabilities.
See: https://pytorch.org/docs/stable/generated/torch.nn.functional.nll_loss.html
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([[-1, 2, -3], [1, -2, 3]])
Y = Tensor([1, 2])
print(t.log_softmax().nll_loss(Y).item())
```
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([[-1, 2, -3], [1, -2, 3]])
Y = Tensor([1, 2])
print(t.log_softmax().nll_loss(Y, reduction='none').numpy())
```
"""
weight = Y.ones_like() if weight is None else weight.gather(0, Y.flatten()).reshape(Y.shape)
masked_weight = weight if ignore_index is None else weight * Y.ne(ignore_index)
nll = -self.gather(1, Y.unsqueeze(1)).squeeze(1) * masked_weight
return nll.sum() / masked_weight.sum() if reduction == "mean" else nll._do_reduction(reduction)
# ***** matrix ops *****
def newton_schulz(self, steps:int, params:tuple[int, ...], eps:float=1.0e-7) -> Self:
tinygrad/nn/onnx.py
+2 -1
View File
@@ -2,7 +2,8 @@
from typing import Any, Sequence, cast, Literal, NamedTuple, Generator
import dataclasses, functools, io, math, types, warnings, pathlib, sys, os, struct, enum
from tinygrad.nn.state import TensorIO
from tinygrad.tensor import Tensor, _broadcast_shape, ReductionStr
from tinygrad.tensor import Tensor, _broadcast_shape
from tinygrad.mixin import ReductionStr
from tinygrad.helpers import getenv, all_same, prod, flatten, make_tuple, argsort, is_numpy_ndarray, get_single_element, polyN
from tinygrad.dtype import DType, ConstType, dtypes, _from_np_dtype, truncate, least_upper_dtype, DTYPES_DICT
from tinygrad.device import is_dtype_supported, Device
tinygrad/tensor.py
+1 -25
View File
@@ -10,7 +10,7 @@ from tinygrad.helpers import argfix, flatten, prod, all_int, round_up, getenv, a
from tinygrad.helpers import resolve_pool_pads, IMAGE, FLOAT16, WINO, Metadata, TRACEMETA, is_numpy_ndarray, TracingKey, cpu_profile
from tinygrad.helpers import suppress_finalizing, disable_gc
from tinygrad.gradient import compute_gradient
from tinygrad.mixin import OpMixin, ReductionStr
from tinygrad.mixin import OpMixin
from tinygrad.uop.ops import smax, UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable, _broadcast_shape
from tinygrad.schedule import ExecItem, create_linear_with_vars, linear_to_schedule
from tinygrad.device import Buffer, canonicalize_device
@@ -1533,30 +1533,6 @@ class Tensor(OpMixin):
qk = qk + attn_mask
return qk.cast(self.dtype).softmax(-1).dropout(dropout_p) @ value
def nll_loss(self, Y:Tensor, weight:Tensor|None=None, ignore_index:int|None=None, reduction:ReductionStr="mean") -> Tensor:
"""
Computes the negative log likelihood loss between log-probabilities and target labels.
NOTE: `self` is log-probabilities and `Y` is the Y labels or class probabilities.
See: https://pytorch.org/docs/stable/generated/torch.nn.functional.nll_loss.html
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([[-1, 2, -3], [1, -2, 3]])
Y = Tensor([1, 2])
print(t.log_softmax().nll_loss(Y).item())
```
```python exec="true" source="above" session="tensor" result="python"
t = Tensor([[-1, 2, -3], [1, -2, 3]])
Y = Tensor([1, 2])
print(t.log_softmax().nll_loss(Y, reduction='none').numpy())
```
"""
weight = Y.ones_like(requires_grad=False) if weight is None else weight[Y]
masked_weight = weight if ignore_index is None else weight * (Y != ignore_index)
nll = -self.gather(1, Y.unsqueeze(1)).squeeze(1) * masked_weight
return nll.sum() / masked_weight.sum() if reduction == "mean" else nll._do_reduction(reduction)
def qr(self) -> tuple[Tensor, Tensor]:
assert self.ndim > 1, f"expected two or more dimensions, got {self.ndim}"
b_shape, m, n = self.shape[:-2], int(self.shape[-2]), int(self.shape[-1])