nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

remove had_counter from rand (#8191)

Browse Source
This commit is contained in:
chenyu
2024-12-12 13:35:39 -05:00
committed by GitHub
parent 2fe98e44cd
commit d586c7e108
1 changed files with 5 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/tensor.py
+5 -8
View File
@@ -492,6 +492,10 @@ class Tensor(SimpleMathTrait):
if device is not None and not isinstance(device, str): raise ValueError(f"rand only supports single device, got {device=}")
_device = device = Device.canonicalize(device)
# if shape has 0, return zero tensor
if (numel := prod(shape)) == 0: return Tensor.zeros(shape, device=_device, dtype=dtype, **kwargs)
num = ceildiv(numel * dtype.itemsize, 4)
# when using MOCKGPU and NV generate rand on CLANG
if getenv("MOCKGPU") and device.startswith("NV"): device = "CLANG"
@@ -501,15 +505,8 @@ class Tensor(SimpleMathTrait):
[int.from_bytes(hashlib.sha256(len(Tensor._device_seeds).to_bytes(4, "big")).digest(), "big"), Tensor._seed],
device=device, dtype=dtypes.uint32, requires_grad=False)
Tensor._device_rng_counters[device] = Tensor([0], device=device, dtype=dtypes.uint32, requires_grad=False)
had_counter = False
else: had_counter = True
# if shape has 0, return zero tensor
if (numel := prod(shape)) == 0: return Tensor.zeros(shape, device=_device, dtype=dtype, **kwargs)
num = ceildiv(numel * dtype.itemsize, 4)
# increment rng counter for devices
if had_counter: Tensor._device_rng_counters[device].assign(Tensor._device_rng_counters[device] + num).contiguous()
else: Tensor._device_rng_counters[device].assign(Tensor._device_rng_counters[device] + num).contiguous()
# threefry random bits
counts0 = (Tensor.arange(ceildiv(num, 2), device=device, dtype=dtypes.uint32, requires_grad=False)+Tensor._device_rng_counters[device])