diff --git a/test/backend/test_multitensor.py b/test/backend/test_multitensor.py index 9e0aa718d6..c8a580fff0 100644 --- a/test/backend/test_multitensor.py +++ b/test/backend/test_multitensor.py @@ -746,6 +746,11 @@ class TestMultiTensor(unittest.TestCase): t2.realize() def test_rand_like_on_shard_axis(self): self.test_rand_like_on_shard(0) + def test_rand_like_on_shard_axis_requires_grad(self): + t = Tensor.empty((16, 16)).shard(devices_2, axis=0) + self.assertIs(t.rand_like(requires_grad=True).requires_grad, True) + self.assertIs(t.rand_like(requires_grad=False).requires_grad, False) + def test_rand_like_from_alu(self): a = Tensor.ones(4, 4).shard(devices_4, axis=0) aa = a + a diff --git a/test/backend/test_tensor.py b/test/backend/test_tensor.py index b98f88540f..95e8061f2a 100644 --- a/test/backend/test_tensor.py +++ b/test/backend/test_tensor.py @@ -260,6 +260,10 @@ class TestTinygrad(unittest.TestCase): b = Tensor.randperm(1000).realize() np.testing.assert_equal(set(b.numpy()), set(range(1000))) + def test_randperm_requires_grad(self): + self.assertIs(Tensor.randperm(5, requires_grad=True).requires_grad, True) + self.assertIs(Tensor.randperm(5, requires_grad=False).requires_grad, False) + def test_randn_isnt_inf_on_zero(self): # simulate failure case of rand handing a zero to randn original_rand, Tensor.rand = Tensor.rand, Tensor.zeros diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 694e571317..f06fe78515 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -618,7 +618,7 @@ class Tensor(OpMixin): if kwargs.get("device") is not None: raise RuntimeError("cannot specify `device` on `*_like` of a multi device tensor") if self.uop.axis is None: return fxn(self.shape, *args, dtype=dtype, **kwargs).shard(self.device) stacked = UOp.mstack(*[fxn(self.uop.shard_shape, *args, device=d, dtype=dtype, **kwargs).uop for d in self.device]) - return Tensor(stacked.multi(self.uop.axis)) + return Tensor(stacked.multi(self.uop.axis), requires_grad=kwargs.get("requires_grad")) def full_like(self, fill_value:ConstType, dtype=None, device=None, requires_grad=None) -> Tensor: """ @@ -816,7 +816,7 @@ class Tensor(OpMixin): print(Tensor.randperm(6).numpy()) ``` """ - return Tensor.rand(n, device=device, **kwargs).argsort().cast(dtype) + return Tensor.rand(n, device=device, **kwargs).argsort().cast(dtype).requires_grad_(kwargs.get("requires_grad")) def multinomial(self:Tensor, num_samples:int = 1, replacement:bool = False) -> Tensor: """