line negative bug fixes

test/backend/test_ops.py

@@ -268,6 +268,11 @@ class TestOps(unittest.TestCase):
 
     self.helper_test_exception([], lambda: torch.meshgrid(x, indexing="bad"), lambda: xt.meshgrid(indexing="bad"), expected=RuntimeError)
 
+  def test_meshgrid_scalar(self):
+    for indexing in ("ij", "xy"):
+      with self.subTest(indexing=indexing):
+        helper_test_op([()], lambda x: torch.meshgrid(x, indexing=indexing)[0], lambda x: x.meshgrid(indexing=indexing)[0])
+
   def test_arange(self):
     helper_test_op([], lambda: torch.arange(10, dtype=torch.int32), lambda: Tensor.arange(10), forward_only=True)
     helper_test_op([], lambda: torch.arange(36, dtype=torch.int32), lambda: Tensor.arange(36), forward_only=True)
@@ -1183,6 +1188,20 @@ class TestOps(unittest.TestCase):
   def test_small_cummax(self):
     helper_test_op([(10)], lambda x: torch.cummax(x, dim=0).values, lambda x: Tensor.cummax(x, axis=0)[0])
     helper_test_op([(10)], lambda x: torch.cummax(x, dim=0).indices.int(), lambda x: Tensor.cummax(x, axis=0)[1], forward_only=True)
+
+  def test_cumextrema_ties(self):
+    for op in ("cummax", "cummin"):
+      for axis in (0, 1, -1):
+        for values in ([[2, 2, 1, 3, 3, 0, 0]] * 2, [[0, 0, 0]] * 2):
+          with self.subTest(op=op, axis=axis, values=values):
+            helper_test_op(None, lambda x: getattr(torch, op)(x, dim=axis).indices.int(),
+                           lambda x: getattr(x, op)(axis)[1], vals=[values], forward_only=True)
+
+  def test_cumextrema_ties_split(self):
+    for op in ("cummax", "cummin"):
+      helper_test_op(None, lambda x: getattr(torch, op)(x, dim=-1).indices.int(), lambda x: getattr(x, op)(-1)[1],
+                     vals=[[[2.0, 2.0, 1.0, 3.0, 3.0, 0.0, 0.0] * 100] * 2], forward_only=True)
+
   @slow_test
   def test_simple_cummax(self):
     helper_test_op([(512)], lambda x: torch.cummax(x, dim=0).values, lambda x: Tensor.cummax(x, axis=0)[0])
@@ -1649,6 +1668,10 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: x.isclose(torch.tensor(1.0)), lambda x: x.isclose(1.0),
                    vals=[[1.0, 1.0 + 1e-7, 2.0, math.inf, -math.inf, math.nan]], forward_only=True)
 
+  def test_isclose_overflow(self):
+    helper_test_op(None, lambda x,y: x.isclose(y, rtol=3),
+                   vals=[[3e38, -3e38, 3e38, 0.0], [-3e38, 3e38, 3e38, 1.0]], forward_only=True)
+
   def test_mean(self):
     helper_test_op([(3,4,5,6)], lambda x: x.mean())
     helper_test_op([()], lambda x: x.mean())
@@ -1693,6 +1716,16 @@ class TestOps(unittest.TestCase):
     helper_test_op([(15, 25, 35)], lambda x: x.var(keepdim=True))
     helper_test_op([(15, 25, 35)], lambda x: x.var(0, keepdim=True, correction=0))
 
+  def test_var_std_integer(self):
+    for op in ("var", "std"):
+      for axis in (None, 0, 1):
+        for correction in (0, 1):
+          for keepdim in (False, True):
+            with self.subTest(op=op, axis=axis, correction=correction, keepdim=keepdim):
+              helper_test_op(None, lambda x: getattr(x.float(), op)(dim=axis, correction=correction, keepdim=keepdim),
+                             lambda x: getattr(x, op)(axis=axis, correction=correction, keepdim=keepdim),
+                             vals=[[[0, 1, 3], [1, 2, 4]]], forward_only=True)
+
   @slow_test
   def test_std(self):
     helper_test_op([(15, 25, 35)], lambda x: x.std())
@@ -1809,6 +1842,21 @@ class TestOps(unittest.TestCase):
     helper_test_op(None, lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(), atol=1e-7, grad_atol=1e-7, vals=[[0.0, 100.0]])
     helper_test_op(None, lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(), vals=[[-math.inf, 0.0, 1.0]], forward_only=True)
 
+  def test_logcumsumexp_scalar_invalid_axis(self):
+    for axis in (-2, 1):
+      with self.subTest(axis=axis):
+        self.helper_test_exception([()], lambda x: torch.logcumsumexp(x, dim=axis), lambda x: x.logcumsumexp(axis), expected=IndexError)
+
+  def test_logcumsumexp_empty(self):
+    for shape, axis in (((0,), 0), ((2, 0, 3), 1), ((2, 0, 3), -1)):
+      with self.subTest(shape=shape, axis=axis):
+        helper_test_op([shape], lambda x: torch.logcumsumexp(x, dim=axis), lambda x: x.logcumsumexp(axis))
+
+  def test_logcumsumexp_nonfinite(self):
+    for values in ([-math.inf, -math.inf], [0., math.inf, -math.inf], [0., math.nan, 1.]):
+      with self.subTest(values=values):
+        helper_test_op(None, lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(), vals=[values], forward_only=True)
+
   def test_sinh(self):
     helper_test_op([(45,65)], lambda x: x.sinh(), grad_atol=1e-6)
     # TODO: backward nan instead of inf
@@ -2189,6 +2237,12 @@ class TestOps(unittest.TestCase):
     helper_test_op([(3,5)], lambda x: x.diagonal(offset=2))  # offset on rectangular
     self.helper_test_exception([(3,3)], lambda x: x.diagonal(dim1=0, dim2=0), expected=RuntimeError)
 
+  def test_diagonal_outside_matrix(self):
+    for shape, dims in (((2, 3), (0, 1)), ((2, 3, 4), (-2, -1)), ((2, 3, 4), (2, 0))):
+      for offset in (-10, -4, 4, 10):
+        with self.subTest(shape=shape, dims=dims, offset=offset):
+          helper_test_op([shape], lambda x: x.diagonal(offset=offset, dim1=dims[0], dim2=dims[1]))
+
   def test_roll(self):
     helper_test_op([(2, 4)], lambda x: x.roll(1))
     helper_test_op([(2, 4)], lambda x: x.roll((1,)))
@@ -3326,6 +3380,16 @@ class TestOps(unittest.TestCase):
       helper_test_op([(12,10)], lambda x: torch.nn.CrossEntropyLoss(label_smoothing=s)(x, torch.tensor(classes)),
                                 lambda x: x.sparse_categorical_crossentropy(Tensor(classes), label_smoothing=s))
 
+  def test_sparse_categorical_crossentropy_default_ignore_index(self):
+    classes = [-1, 0, 2, -1]
+    for reduction in ("none", "sum", "mean"):
+      for smoothing in (0.0, 0.3, 1.0):
+        with self.subTest(reduction=reduction, smoothing=smoothing):
+          helper_test_op([(4, 3)],
+                         lambda x: torch.nn.functional.cross_entropy(x, torch.tensor(classes), ignore_index=-1,
+                                                                     reduction=reduction, label_smoothing=smoothing),
+                         lambda x: x.sparse_categorical_crossentropy(Tensor(classes), reduction=reduction, label_smoothing=smoothing))
+
   def test_nll_loss(self):
     target = np.random.randint(0, 10, (32,), dtype=np.int32).tolist()
     helper_test_op([(32,10)],
@@ -3439,6 +3503,31 @@ class TestOps(unittest.TestCase):
     if not COMPILE_ONLY: assert t == -1
 
 class TestOpsUint8(unittest.TestCase):
+  def test_lerp_integer_end(self):
+    for dtype in dtypes.ints:
+      with self.subTest(dtype=dtype):
+        actual = Tensor([[10], [100]], dtype=dtypes.uint8).lerp(Tensor([20, 20, 100], dtype=dtype), Tensor([0., 0.5, 1.]))
+        self.assertEqual(actual.dtype, dtypes.uint8)
+        actual.realize()
+        if not COMPILE_ONLY: np.testing.assert_equal(actual.numpy(), [[10, 15, 100], [100, 60, 100]])
+
+  def test_lerp_float_end(self):
+    helper_test_op(None, lambda x,y,w: x.float().lerp(y, w), lambda x,y,w: x.cast(dtypes.uint8).lerp(y, w),
+                   vals=[[[10], [100]], [20.5, 9.5, -5.5], [0., 0.5, 1.]], forward_only=True)
+
+  def test_interpolate_bilinear_full_range(self):
+    for values in ([[0, 255]], [[255, 0]], [[1, 200]], [[0, 255], [255, 0]]):
+      for size in ((1, 3), (5, 10)):
+        for align_corners in (False, True):
+          with self.subTest(values=values, size=size, align_corners=align_corners):
+            image = torch.tensor([[values]], dtype=torch.uint8)
+            expected = torch.nn.functional.interpolate(image, size=size, mode="bilinear", align_corners=align_corners)
+            actual = Tensor(image.numpy()).interpolate(size, align_corners=align_corners)
+            self.assertEqual(actual.dtype, dtypes.uint8)
+            # Midpoints are exact; other weights can differ by one with 7-bit fixed-point coefficients.
+            actual.realize()
+            if not COMPILE_ONLY: np.testing.assert_allclose(actual.numpy(), expected.numpy(), rtol=0, atol=0 if size == (1, 3) else 1)
+
   def test_cast(self):
     helper_test_op([(2,3,64,64)], lambda x: x.type(torch.uint8), lambda x: x.cast('uint8'), forward_only=True, low=0, high=255)
 

test/unit/test_dtype_spec.py

@@ -151,6 +151,17 @@ class TestTypeSpec(unittest.TestCase):
     _assert_eq(Tensor.arange(5.0, 3.0), dtypes.default_float, np.arange(5.0, 3.0))
 
 class TestAutoCastType(unittest.TestCase):
+  @unittest.skipUnless(dtypes.float64 in supported_dtypes, "need float64")
+  def test_linspace_float64_precision(self):
+    for start, stop in ((1., 1.+1e-8), (1.+1e-8, 1.), (1e10, 1e10+1)):
+      with self.subTest(start=start, stop=stop):
+        out = Tensor.linspace(start, stop, 3, dtype=dtypes.float64)
+        self.assertEqual(out.dtype, dtypes.float64)
+        np.testing.assert_allclose(out.numpy(), np.linspace(start, stop, 3), rtol=1e-15, atol=0)
+    with Context(DEFAULT_FLOAT=dtypes.float64):
+      out = Tensor.linspace(10**10, 10**10+2, 3, dtype=dtypes.int64)
+      np.testing.assert_array_equal(out.numpy(), [10**10, 10**10+1, 10**10+2])
+
   def test_int_sqrt(self):
     _assert_eq(Tensor([1, 4, 9, 16]).sqrt(), dtypes.default_float, [1, 2, 3, 4])
 
@@ -222,6 +233,13 @@ class TestAutoCastType(unittest.TestCase):
     t.square().mean().backward()
     np.testing.assert_allclose(t.grad.numpy().flatten(), [60000 * 2 / (N*N)] * N*N)
 
+  def test_var_integer_fractional(self):
+    for dtype in [*dtype_ints, dtypes.bool]:
+      with self.subTest(dtype=dtype):
+        out = Tensor([0, 1], dtype=dtype).var()
+        self.assertEqual(out.dtype, dtypes.float32)
+        np.testing.assert_allclose(out.numpy(), 0.5)
+
   @unittest.skipUnless(dtypes.half in supported_dtypes, "need half")
   def test_var_half_precision_large_n(self):
     # the element count (70000) exceeds half max (65504): the denominator must not be materialized in half

test/unit/test_state.py

@@ -0,0 +1,82 @@
+import unittest
+from collections import OrderedDict, namedtuple
+from types import SimpleNamespace
+from tinygrad import Tensor
+from tinygrad.nn.state import get_parameters, get_state_dict, load_state_dict
+
+
+class TestStateDict(unittest.TestCase):
+  def test_container_subclasses(self):
+    class TensorDict(dict): pass
+    class TensorList(list): pass
+    class TensorTuple(tuple): pass
+    weight = Tensor([1., 2.])
+    for container, key in ((TensorDict(weight=weight), "weight"), (TensorList([weight]), "0"), (TensorTuple([weight]), "0")):
+      with self.subTest(container=type(container).__name__):
+        container.description = "model weights"
+        model = SimpleNamespace(layers=container)
+        state = get_state_dict(model)
+        self.assertEqual(list(state), [f"layers.{key}"])
+        self.assertIs(state[f"layers.{key}"], weight)
+        params = get_parameters(model)
+        self.assertEqual(len(params), 1)
+        self.assertIs(params[0], weight)
+
+  def test_namedtuple_and_ordered_dict(self):
+    first, second = Tensor([1.]), Tensor([2.])
+    pair = namedtuple("Pair", ["first", "second"])(first, second)
+    state = get_state_dict(OrderedDict(pair=pair))
+    self.assertEqual(list(state), ["pair.first", "pair.second"])
+    self.assertIs(state["pair.first"], first)
+    self.assertIs(state["pair.second"], second)
+
+  def test_load_container_subclass(self):
+    class TensorDict(dict): pass
+    weight = Tensor([1., 2.])
+    model = TensorDict(weight=weight)
+    loaded = load_state_dict(model, {"weight": Tensor([3., 4.])}, verbose=False)
+    self.assertEqual(len(loaded), 1)
+    self.assertIs(loaded[0], weight)
+    self.assertEqual(weight.tolist(), [3., 4.])
+
+  def test_container_tensor_attributes(self):
+    class TensorDict(dict): pass
+    class TensorList(list): pass
+    class TensorTuple(tuple): pass
+    for container_type in (TensorDict, TensorList, TensorTuple):
+      with self.subTest(container=container_type.__name__):
+        model = container_type()
+        model.weight = Tensor([1., 2.])
+        state = get_state_dict(model)
+        self.assertEqual(list(state), ["weight"])
+        self.assertIs(state["weight"], model.weight)
+        params = get_parameters(model)
+        self.assertEqual(len(params), 1)
+        self.assertIs(params[0], model.weight)
+        loaded = load_state_dict(model, {"weight": Tensor([3., 4.])}, verbose=False)
+        self.assertEqual(len(loaded), 1)
+        self.assertIs(loaded[0], model.weight)
+        self.assertEqual(model.weight.tolist(), [3., 4.])
+
+  def test_container_contents_and_attributes(self):
+    class TensorDict(dict): pass
+    class TensorList(list): pass
+    class TensorTuple(tuple): pass
+    item, weight = Tensor([1.]), Tensor([2.])
+    for model, key in ((TensorDict(item=item), "item"), (TensorList([item]), "0"), (TensorTuple([item]), "0")):
+      with self.subTest(container=type(model).__name__):
+        model.weight = weight
+        state = get_state_dict(model, prefix="model.")
+        self.assertEqual(list(state), [f"model.{key}", "model.weight"])
+        self.assertIs(state[f"model.{key}"], item)
+        self.assertIs(state["model.weight"], weight)
+
+  def test_container_attribute_precedence(self):
+    class TensorDict(dict): pass
+    model = TensorDict(weight=Tensor([1.]))
+    model.weight = Tensor([2.])
+    self.assertIs(get_state_dict(model)["weight"], model.weight)
+
+
+if __name__ == '__main__':
+  unittest.main()

tinygrad/mixin/elementwise.py

@@ -647,10 +647,9 @@ class ElementwiseMixin(CreationMixin):
     ```
     """
     other = self.ufix(other)
-    is_finite_close = self.isfinite() & other.isfinite() & ((self - other).abs() <= atol + rtol * other.abs())
-    is_infinite_close = (self.isinf() | other.isinf()) & self.eq(other)
-    is_nan_close = (self.isnan() & other.isnan()) & equal_nan
-    return is_finite_close | is_infinite_close | is_nan_close
+    error = (self - other).abs()
+    is_finite_close = error.isfinite() & (error <= atol + rtol * other.abs())
+    return self.eq(other) | is_finite_close | (self.isnan() & other.isnan() & equal_nan)
 
   def ceil(self) -> Self:
     """
@@ -1087,7 +1086,7 @@ class ElementwiseMixin(CreationMixin):
     print(Tensor([1., 2., 3.]).lerp(Tensor([4., 5., 6.]), 0.5).numpy())
     ```
     """
-    if self.dtype == dtypes.uint8 and not isinstance(weight, ConstType):
-      w_i = (weight * (1<<(W_PREC:=7)) + 0.5).cast(dtypes.int16)
-      return (self+(((end - self).cast(dtypes.int8) * w_i + (1<<W_PREC-1)).cast(dtypes.uint16) >> W_PREC)).cast(dtypes.uint8)
+    if self.dtype == dtypes.uint8 and not end.is_floating_point() and not isinstance(weight, ConstType):
+      weight_int = (weight * 128 + 0.5).cast(dtypes.int32)  # 7 fractional bits
+      return ((self * (128 - weight_int) + end.cast(dtypes.int32) * weight_int + 64) >> 7).cast(dtypes.uint8)
     return self + (end - self) * weight

tinygrad/mixin/movement.py

@@ -493,9 +493,8 @@ class MovementMixin:
     ```
     """
     if indexing not in ("ij", "xy"): raise RuntimeError(f'indexing must be in ("ij", "xy"), got {indexing}')
-    if len(tensors:=(self, *args)) == 1: return tensors
-    basis = tuple(range(len(tensors))) if indexing == "ij" else (1, 0) + tuple(range(2, len(tensors)))
-    tensors = tuple(t.reshape((-1,) + (1,)*(len(args) - i)) for i,t in zip(basis, tensors))
+    basis = tuple(range(len(args)+1)) if indexing == "ij" or not args else (1, 0) + tuple(range(2, len(args)+1))
+    tensors = tuple(t.reshape((-1,) + (1,)*(len(args) - i)) for i,t in zip(basis, (self, *args)))
     output_shape = _broadcast_shape(*(t.shape for t in tensors))
     return tuple(t._broadcast_to(output_shape) for t in tensors)
 
@@ -528,8 +527,8 @@ class MovementMixin:
     """
     if (dim1:=self._resolve_dim(dim1)) == (dim2:=self._resolve_dim(dim2)): raise RuntimeError("dim1 and dim2 cannot be the same dimension")
     x = self.permute(*[i for i in range(self.ndim) if i != dim1 and i != dim2], dim1, dim2)
-    if offset >= 0: x = x.shrink(tuple(None for _ in x.shape[:-1]) + ((offset, x.shape[-1]),))
-    else: x = x.shrink(tuple(None for _ in x.shape[:-2]) + ((-offset, x.shape[-2]), None))
+    if offset >= 0: x = x.shrink((None,)*(x.ndim-1) + ((min(offset, x.shape[-1]), x.shape[-1]),))
+    else: x = x.shrink((None,)*(x.ndim-2) + ((min(-offset, x.shape[-2]), x.shape[-2]), None))
     if (d := min(int(x.shape[-2]), int(x.shape[-1]))) <= 0: return x.reshape(*x.shape[:-2], 0)
     nones, x = tuple(None for _ in x.shape[:-2]), x.shrink_to(tuple(None for _ in x.shape[:-2]) + (d, d))
     return x.flatten(-2).pad_to(nones+(d*(d+1),)).unflatten(-1, (d, d+1)).shrink_to(nones+(None, 1)).squeeze(-1)

tinygrad/mixin/op.py

@@ -206,7 +206,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     if steps < 0: raise ValueError("number of steps must be non-negative")
     if (dtype := to_dtype(dtype or dtypes.default_float)) == dtypes.bool: raise ValueError("linspace with bool dtype is not supported")
     if steps == 1: return cls.full((1,), start, dtype=dtype, buffer=False)
-    return (start + cls.arange(steps, dtype=dtypes.default_float) * ((stop - start) / (steps - 1))).cast(dtype)
+    return (start + cls.arange(steps, dtype=least_upper_dtype(dtype, dtypes.default_float)) * ((stop - start) / (steps - 1))).cast(dtype)
 
   @classmethod
   def eye(cls, n:int, m:int|None=None, dtype:DTypeLike|None=None) -> Self:
@@ -542,11 +542,10 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     print(t.var(axis=1).numpy())
     ```
     """
-    output_dtype = self.dtype if dtypes.is_float(self.dtype) else dtypes.float32
     squares = (self - self.mean(axis=axis, keepdim=True)).square()
     n = prod([si for si, so in zip(self.shape, squares.sum(axis=axis, keepdim=True).shape) if resolve(si != so)])
-    numerator = squares.cast(sum_acc_dtype(self.commit_dtype())).sum(axis=axis, keepdim=keepdim)
-    return numerator.div(smax(n - correction, 0)).cast(output_dtype)
+    numerator = squares.sum(axis=axis, keepdim=keepdim, dtype=sum_acc_dtype(squares.commit_dtype()))
+    return numerator.div(smax(n - correction, 0)).cast(squares.dtype)
 
   def var_mean(self, axis:int|Sequence[int]|None=None, keepdim=False, correction=1) -> tuple[Self, Self]:
     """
@@ -807,11 +806,10 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     ```
     """
     if self.ndim == 0: return self._split_cumalu(axis, Ops.MAX), type(self).zeros(self.shape, dtype=dtypes.int32, buffer=False)
-    values, n = self._split_cumalu(axis, Ops.MAX), int(self.shape[axis])
-    x, values_t = self.transpose(axis, -1), values.transpose(axis, -1)
-    match = x.unsqueeze(-1).eq(values_t.unsqueeze(-2)) * self._tri(n, n)
-    idx = (-(match * type(self).arange(n, 0, -1).reshape(n, 1)).max(-2) + n).cast(dtypes.int32)
-    return values, idx.transpose(-1, axis)
+    values = self._split_cumalu(axis, Ops.MAX)
+    # Record the latest index matching the running maximum, then carry it forward.
+    idx = self.eq(values).transpose(axis, -1) * type(self).arange(self.shape[axis], dtype=dtypes.int32)
+    return values, idx._split_cumalu(-1, Ops.MAX).transpose(-1, axis)
 
   def cummin(self, axis:int=0) -> tuple[Self, Self]:
     """
@@ -851,14 +849,12 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     print(t.logcumsumexp(axis=1).numpy())
     ```
     """
+    axis = self._resolve_dim(axis)
     if self.ndim == 0: return self
     x = self.transpose(axis, -1)
-    last_dim_size = x.shape[-1]
-    x_unsqueezed = x.unsqueeze(-2)
-    x_cummax = (mx:=x.cummax(-1)[0].detach()).isfinite().where(mx, 0)
-    mask = self._tri(last_dim_size, last_dim_size, 1).logical_not()
-    ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), self.dtype.min).exp().sum(-1).log() + x_cummax
-    return ret.transpose(-1, axis)
+    mask = self._tri(x.shape[-1], x.shape[-1], 1)
+    prefixes = mask.where(-math.inf, x.unsqueeze(-2))
+    return prefixes.logsumexp(-1).transpose(-1, axis)
 
   def argmax(self, axis=None, keepdim=False) -> Self:
     """
@@ -1738,10 +1734,10 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     if Y.device is not None and self.device is not None and Y.device != self.device:
       raise RuntimeError(f"expected Y and self on the same device, {Y.device=}, {self.device=}")
     log_probs = self.log_softmax()
-    loss_mask = Y.ne(ignore_index) if ignore_index != -1 else Y.const_like(True, dtypes.bool)
-    y = Y.unsqueeze(-1)._one_hot_along_dim(self.shape[-1], dim=-1) * loss_mask.unsqueeze(-1)
-    smoothing = label_smoothing * (log_probs.mean(-1) * loss_mask)
-    unreduced = ((1 - label_smoothing) * (log_probs * y).sum(-1) + smoothing)
+    loss_mask = Y.ne(ignore_index)
+    y = Y.unsqueeze(-1)._one_hot_along_dim(self.shape[-1], dim=-1)
+    smoothing = label_smoothing * log_probs.mean(-1)
+    unreduced = ((1 - label_smoothing) * (log_probs * y).sum(-1) + smoothing) * loss_mask
     return -unreduced.sum() / loss_mask.sum() if reduction == "mean" else -unreduced._do_reduction(reduction)
 
   def cross_entropy(self, Y:Self, reduction:ReductionStr="mean", label_smoothing:float=0.0) -> Self:

tinygrad/nn/state.py

@@ -1,5 +1,4 @@
 import json, pathlib, struct, functools, io, zlib
-from collections import OrderedDict
 from typing import Any, Callable, BinaryIO, Iterable, cast
 from tinygrad.tensor import Tensor
 from tinygrad.dtype import dtypes
@@ -102,13 +101,12 @@ def get_state_dict(obj, prefix:str='', tensor_type=Tensor) -> dict[str, Tensor]:
   """
   if isinstance(obj, tensor_type): return {prefix.strip('.'):obj}
   if hasattr(obj, '_asdict'): return get_state_dict(obj._asdict(), prefix, tensor_type)  # namedtuple
-  if isinstance(obj, OrderedDict): return get_state_dict(dict(obj), prefix, tensor_type)
-  if hasattr(obj, '__dict__'): return get_state_dict(obj.__dict__, prefix, tensor_type)
   state_dict = {}
   if isinstance(obj, (list, tuple)):
     for i,x in enumerate(obj): state_dict.update(get_state_dict(x, f"{prefix}{str(i)}.", tensor_type))
   elif isinstance(obj, dict):
     for k,v in obj.items(): state_dict.update(get_state_dict(v, f"{prefix}{str(k)}.", tensor_type))
+  if hasattr(obj, '__dict__'): state_dict.update(get_state_dict(obj.__dict__, prefix, tensor_type))
   return state_dict
 
 def get_parameters(obj) -> list[Tensor]: