tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 23:06:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

minor rand cleanups (#15982)

Browse Source
This commit is contained in:
chenyu
2026-04-29 22:22:29 -04:00
committed by GitHub
parent 53a28bafbd
commit b73248958a
1 changed files with 5 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/tensor.py
+5 -6
View File
@@ -554,16 +554,15 @@ class Tensor(OpMixin):
@staticmethod
def _next_counter(device:str, num:int) -> tuple[Tensor, Tensor]:
if device not in Tensor._device_seeds:
Tensor._device_seeds[device] = Tensor(
[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, 0], device=device, dtype=dtypes.uint32, requires_grad=False).contiguous()
seed = [int.from_bytes(hashlib.sha256(len(Tensor._device_seeds).to_bytes(4, "big")).digest(), "big"), Tensor._seed]
Tensor._device_seeds[device] = Tensor(seed, device=device, dtype=dtypes.uint32, requires_grad=False)
Tensor._device_rng_counters[device] = Tensor([0, 0], device=device, dtype=dtypes.uint32, requires_grad=False)
counter = Tensor._device_rng_counters[device]
new_low = counter[0:1] + (num & 0xffffffff)
new_high = counter[1:2] + (num >> 32) + (new_low < counter[0]).cast(dtypes.uint32)
new_high = counter[1:2] + (num >> 32) + (new_low < counter[0])
counter.assign(new_low.cat(new_high))
low = counter[0:1] - (num & 0xffffffff)
high = counter[1:2] - (num >> 32) - (counter[0] < (num & 0xffffffff)).cast(dtypes.uint32)
high = counter[1:2] - (num >> 32) - (counter[0] < (num & 0xffffffff))
return Tensor._device_seeds[device], low.cat(high)
@staticmethod