fix dtype mismatch for bool ops in multi (tinygrad#5299)

ignaciosica · Jul 6, 2024 · c1e166c · c1e166c
1 parent fc03fc0
commit c1e166c
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 2 deletions.
diff --git a/test/test_multitensor.py b/test/test_multitensor.py
@@ -135,6 +135,15 @@ def test_four_add(self):
     O = X + W
     np.testing.assert_allclose(O.numpy(), 2)
 
+  def test_elementwise_dtype(self):
+    Tensor.manual_seed(0)
+    X = Tensor.randn(8, 8).realize()
+    W = Tensor.randn(8, 8).realize()
+    X.shard_(devices_4, 0)
+    W.shard_(devices_4, 0)
+    O = X.shrink(((0, 2), None)) * W.shrink(((0, 2), None)) < 2
+    np.testing.assert_allclose(O.numpy(), X.numpy()[0:2]*W.numpy()[0:2] < 2)
+
   @given(strat.sampled_from((4, 5)), strat.sampled_from((devices_2, devices_3)), strat.sampled_from((ReduceOps.SUM, ReduceOps.MAX)),
          strat.sampled_from((None, 0, 1)), strat.sampled_from((None, 0, 1)), strat.sampled_from((1, 0, -1)))
   def test_simple_reduce(self, N, devices, rop, shard_axis, reduce_axis, sign):

diff --git a/tinygrad/multi.py b/tinygrad/multi.py
@@ -112,8 +112,11 @@ def e(self, op:Union[LoadOps, UnaryOps, BinaryOps, TernaryOps], *in_srcs:MultiLa
       if mlb.axis == axis or not_all_real: srcs.append(mlb.lbs)
       elif mlb.axis is None and axis is not None: srcs.append(to_sharded(mlb.lbs, axis))
       else: srcs.append(to_sharded([mlb.copy_to_device(lb.device) for lb in mlb.lbs], axis))
-    # NOTE: lsrcs[-1].const(0) is correct for where
-    return MultiLazyBuffer([lsrcs[0].e(op, *lsrcs[1:], arg=arg) if r else lsrcs[-1].const(0) for lsrcs,r in zip(zip(*srcs),new_real)], axis, new_real)
+    new_real_lbs = {i:lsrcs[0].e(op, *lsrcs[1:], arg=arg) for i,(lsrcs,r) in enumerate(zip(zip(*srcs),new_real)) if r}
+    # NOTE: const dtype should match real
+    real_dtype = next(iter(new_real_lbs.values())).dtype
+    return MultiLazyBuffer([new_real_lbs[i] if r else lsrcs[0].const(0).cast(real_dtype) \
+        for i, (lsrcs,r) in enumerate(zip(zip(*srcs),new_real))], axis, new_real)
 
   def _shape_to_single_shard(self, shape:Tuple[sint, ...], lb:LazyBuffer) -> Tuple[sint, ...]:
     return tuple(lb.shape[self.axis] if a == self.axis else s for a,s in enumerate(shape))