fix shape broadcast for symbolic (#16970)

* fix shape broadcast for symbolic

0 or smax was wrong for the unresolved case

* test with null

test/backend/test_tensor_variable.py

@@ -49,11 +49,9 @@ class TestTensorVariable(unittest.TestCase):
     # NOTE: the buffer dim must cover the variable's vmax
     vv = Variable("a", 1, 10).bind(2)
     self.assertEqual((Tensor.ones(10).contiguous()[:vv] * Tensor(vv)).sum().item(), 4.0)
+    # a vmin=0 symbolic dim broadcasts too
     v0 = Variable("z", 0, 10).bind(2)
-    # TODO: broadcasting a vmin=0 symbolic dim fails, max(dim, 1) cannot be proven equal to dim
-    try:
-      self.assertEqual((Tensor.ones(10).contiguous()[:v0] * Tensor(v0)).sum().item(), 4.0)
-    except IndexError: pass
+    self.assertEqual((Tensor.ones(10).contiguous()[:v0] * Tensor(v0)).sum().item(), 4.0)
 
   def test_inner_tvar_node(self):
     vv = Variable("w", 0, 10).bind(2)

test/null/test_symbolic_tensor.py

@@ -1,6 +1,23 @@
 import unittest
 from tinygrad import Variable
 from tinygrad.tensor import Tensor
+from tinygrad.uop.ops import _broadcast_shape
+
+class TestBroadcastShape(unittest.TestCase):
+  def test_symbolic(self):
+    v = Variable("v", 1, 10)
+    self.assertEqual(_broadcast_shape((v,), (1,)), (v,))
+    self.assertEqual(_broadcast_shape((v,), ()), (v,))
+    self.assertEqual(_broadcast_shape((v,), (v,)), (v,))
+    with self.assertRaises(IndexError): _broadcast_shape((v,), (5,))
+
+  def test_symbolic_vmin_zero(self):
+    # a symbolic dim that may be 0 still broadcasts against 1 to itself
+    v0 = Variable("v0", 0, 10)
+    self.assertEqual(_broadcast_shape((v0,), (1,)), (v0,))
+    self.assertEqual(_broadcast_shape((v0,), ()), (v0,))
+    self.assertEqual(_broadcast_shape((3, v0), (3, 1)), (3, v0))
+    with self.assertRaises(IndexError): _broadcast_shape((v0,), (5,))
 
 class TestSymbolic(unittest.TestCase):
   def assert_tuple_equal(self, x, y):

tinygrad/uop/ops.py

@@ -65,11 +65,13 @@ def _align_left(*shapes:tuple[sint, ...]) -> tuple[tuple[sint, ...], ...]:
   max_dim = max(len(s) for s in shapes)
   return tuple((1,)*(max_dim-len(s))+s for s in shapes)
 def _broadcast_shape(*shapes:tuple[sint, ...]) -> tuple[sint, ...]:
-  shaped_aligned_left = _align_left(*shapes)
-  ret = tuple(0 if 0 in nth_dim_sizes else smax(nth_dim_sizes) for nth_dim_sizes in zip(*shaped_aligned_left))
-  if not all(resolve(s == ns) or resolve(s == 1) for shape in shaped_aligned_left for s,ns in zip(shape, ret)):
-    raise IndexError(f"shape mismatch: objects cannot be broadcast to a single shape {shapes}")
-  return ret
+  # per right-aligned dim: sizes of 1 broadcast to the others, which must all agree
+  ret = []
+  for sizes in zip(*_align_left(*shapes)):
+    if len(rest:=dedup([s for s in sizes if isinstance(s, UOp) or s != 1])) > 1:
+      raise IndexError(f"shape mismatch: objects cannot be broadcast to a single shape {shapes}")
+    ret.append(rest[0] if rest else 1)
+  return tuple(ret)
 
 def ssimplify(uop:sint): return uop.ssimplify() if isinstance(uop, UOp) else uop
 def sym_infer(uop: UOp|int, var_vals: dict[str, int]) -> int: return uop.sym_infer(var_vals) if isinstance(uop, UOp) else uop