bring back old sharded rand behavior (#10842)

test/test_multitensor.py

@@ -681,13 +681,15 @@ class TestMultiTensor(unittest.TestCase):
     self.assertEqual(d0_rand, d1_rand_flip)
     self.assertEqual(d1_rand, d0_rand_flip)
 
-  def test_rand_like_on_shard(self):
-    t = Tensor.empty((16, 16)).shard(devices_2)
+  def test_rand_like_on_shard(self, axis=None):
+    t = Tensor.empty((16, 16)).shard(devices_2, axis=axis)
     t2 = Tensor.rand_like(t)
     self.assertEqual(t.shape, t2.shape)
     self.assertEqual(t.device, t2.device)
     self.assertEqual(t.dtype, t2.dtype)
     self.assertEqual(t.uop.axis, t2.uop.axis)
+    t2.realize()
+  def test_rand_like_on_shard_axis(self): self.test_rand_like_on_shard(0)
 
   def test_rand_like_from_alu(self):
     a = Tensor.ones(4, 4).shard(devices_4, axis=0)

tinygrad/tensor.py

@@ -722,7 +722,12 @@ class Tensor(MathTrait):
     dtype = kwargs.pop("dtype", self.dtype)
     if isinstance(self.device, tuple):
       if kwargs.get("device") is not None: raise RuntimeError("cannot specify `device` on `rand_like` of a multi device tensor")
-      return Tensor.rand(*self.shape, dtype=dtype, **kwargs).shard(self.device, self.uop.axis)
+      if self.uop.axis is None: return Tensor.rand(*self.shape, dtype=dtype, **kwargs).shard(self.device)
+      contiguous = kwargs.pop("contiguous", True)
+      sharded_shape = tuple(s//len(self.device) if a==self.uop.axis else s for a,s in enumerate(self.shape))
+      rands = UOp(Ops.MSTACK, dtype=dtype,
+                  src=tuple([Tensor.rand(sharded_shape, device=d, dtype=dtype, contiguous=contiguous, **kwargs).uop for d in self.device]))
+      return Tensor(UOp.multi(rands, axis=self.uop.axis), device=self.device, dtype=dtype, **kwargs)
     return Tensor.rand(*self.shape, device=kwargs.pop("device", self.device), dtype=dtype, **kwargs)
 
   # ***** rng hlops *****