fix cast to weak twice [pr] (#17548)

also no gradient for weak target

test/unit/test_dtype_weak.py

@@ -93,6 +93,13 @@ class TestWeakPromotion(unittest.TestCase):
       out = Tensor(1.0, dtype=dtypes.float32, device="CPU") / denom
       self.assertAlmostEqual(out.item(), 1 / (70000 + 1e-5), places=10)
 
+  def test_stacked_weak_casts_convert_each_kind(self):
+    # each weak cast is a kind conversion: weakint truncates before weakfloat re-lifts (neither is only a marker)
+    x = Tensor([2.5, -3.7], dtype=dtypes.float32, device="CPU")
+    stacked = x.cast(dtypes.weakint).cast(dtypes.weakfloat)
+    self.assertIs(stacked.dtype, dtypes.weakfloat)
+    self.assertEqual(stacked.tolist(), [2.0, -3.0])
+
   def test_uop_scalar_const_lifts_kind(self):
     for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint),
                                                  (dtypes.int32, 1, dtypes.int32, dtypes.weakint),

test/unit/test_gradient.py

@@ -51,6 +51,10 @@ class TestTensorGradient(unittest.TestCase):
     with self.assertRaises(RuntimeError): x.sum().gradient(x)
     with self.assertRaises(RuntimeError): x.float().sum().gradient(x)
 
+  def test_const_target_raise(self):
+    t = Tensor(2.0)
+    with self.assertRaises(RuntimeError): (t * 2.0).gradient(t)
+
   def test_copy_to_device_gradient(self):
     t = Tensor([1.0, 2, 3]).realize()
     t.to("CPU:1").square().sum().backward()
@@ -100,7 +104,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_where_gradient(self):
     # WHERE with a bare ()-shape branch: the scalar's gradient counts the positions where it is selected
-    cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0)
+    cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0, dtype=dtypes.float32)
     dw = Tensor(cond.uop.alu(Ops.WHERE, x.uop, w.uop)).sum().gradient(w)[0]
     self.assertEqual(dw.shape, ())
     self.assertEqual(dw.item(), 1.0)
@@ -109,7 +113,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_alu_gradient(self):
     # MUL with a bare ()-shape src, no EXPAND in the graph
-    x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0)
+    x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0, dtype=dtypes.float32)
     m = x.uop.alu(Ops.MUL, w.uop)
     self.assertIs(m.src[1], w.uop)
     dw = Tensor(m).sum().gradient(w)[0]
@@ -118,7 +122,7 @@ class TestTensorGradient(unittest.TestCase):
 
   def test_implicit_broadcast_intermediate_accumulation(self):
     # s is used directly and through an implicit broadcast edge, each edge's gradient reduces to s's shape before they sum
-    x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5)
+    x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5, dtype=dtypes.float32)
     s = p.sin()
     z = Tensor(x.uop.alu(Ops.MUL, s.uop)).sum() + s
     dp = z.gradient(p)[0]

tinygrad/mixin/op.py

@@ -460,6 +460,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     """
     assert gradient is not None or self.shape == tuple(), "when no gradient is provided, backward must be called on a scalar tensor"
     if not (self.is_floating_point() and all(t.is_floating_point() for t in targets)): raise RuntimeError("only float Tensors have gradient")
+    if any(t.dtype in dtypes.weaks for t in targets): raise RuntimeError("cannot take gradient wrt a weak Tensor")
     from tinygrad.mixin.gradient import compute_gradient
     if gradient is None: gradient = self.const_like(1.0)
     target_uops = [t._uop for t in targets]

tinygrad/uop/weak.py

@@ -16,10 +16,10 @@ def lower_weak_node(u:UOp) -> UOp|None:
 
 pm_lower_weak = PatternMatcher([
   (UPat(Ops.CONST, dtype=dtypes.weaks, name="u"), lambda u: UOp.const(u.val, select_dtype(u)).cast(u.dtype)),
-  # two stacked weak casts are a weakint value used as weakfloat (or vice versa): resolve the inner one at the outer kind's default.
+  # two stacked weak casts are two kind conversions: each resolves at its own kind's default
   # a SINGLE weak cast is never rewritten here, each consumer absorbs it on its own edge (see lower_weak_srcs)
   (UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat(Ops.CAST, dtype=dtypes.weaks, src=(UPat.var("x"),)),), name="u"),
-   lambda u,x: x.cast(select_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None),
+   lambda u,x: x.cast(select_dtype(u.src[0])).cast(select_dtype(u)).cast(u.dtype) if x.dtype not in dtypes.weaks else None),
   # Binary can widen from the bounds, all other nodes derive from the lowered sources.
   # a weakfloat Unary (sin/exp2/...) must resolve here, before the transcendental decomposition
   (UPat(GroupOp.Binary|GroupOp.Unary|{Ops.WHERE, Ops.RANGE, Ops.STACK, Ops.SPECIAL}, name="u"), lower_weak_node),