diff --git a/test/test_multitensor.py b/test/test_multitensor.py index 0f6f16b48a..7e218b8bcb 100644 --- a/test/test_multitensor.py +++ b/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) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 8ce86cb012..aa3ef8ef73 100644 --- a/tinygrad/tensor.py +++ b/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 *****