no weak dtype for randn_like either (#17055)

This commit is contained in:
chenyu
2026-07-16 18:29:12 -04:00
committed by GitHub
parent 3bfd62e915
commit 88826a6f35
2 changed files with 3 additions and 1 deletions
+1
View File
@@ -11,6 +11,7 @@ class TestWeakPromotion(unittest.TestCase):
def test_rand_requires_concrete(self):
with self.assertRaises(ValueError): Tensor.rand(2, dtype=dtypes.weakfloat)
with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).rand_like()
with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).randn_like()
def test_sum_stays_weak(self):
for weak, value in ((dtypes.weakint, 1), (dtypes.weakfloat, 1.0)):
+2 -1
View File
@@ -97,9 +97,10 @@ class RandMixin(OpMixin):
print(Tensor.randn_like(t).numpy())
```
"""
if (dt:=to_dtype(dtype or self.dtype)) in dtypes.weaks and dtype is None: raise ValueError(f"randn_like requires an explicit dtype for {dt}")
src = self.stack(self).rand_like(**{**kwargs, "dtype": dtypes.float32})
# https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform
return src[0].mul(2*math.pi).cos().mul((1 - src[1]).log().mul(-2).sqrt()).cast(to_dtype(dtype or self.dtype))
return src[0].mul(2*math.pi).cos().mul((1 - src[1]).log().mul(-2).sqrt()).cast(dt)
@classmethod
def randn(cls, *shape, dtype:DTypeLike|None=None, **kwargs) -> Self: