MCPcopy Create free account
hub / github.com/tinygrad/tinygrad / _next_counter

Method _next_counter

tinygrad/tensor.py:553–564  ·  view source on GitHub ↗
(device:str, num:int)

Source from the content-addressed store, hash-verified

551
552 @staticmethod
553 def _next_counter(device:str, num:int) -> tuple[Tensor, Tensor]:
554 if device not in Tensor._device_seeds:
555 seed = [int.from_bytes(hashlib.sha256(len(Tensor._device_seeds).to_bytes(4, "big")).digest(), "big"), Tensor._seed]
556 Tensor._device_seeds[device] = Tensor(seed, device=device, dtype=dtypes.uint32)
557 Tensor._device_rng_counters[device] = Tensor([0, 0], device=device, dtype=dtypes.uint32)
558 counter = Tensor._device_rng_counters[device]
559 new_low = counter[0:1] + (num & 0xffffffff)
560 new_high = counter[1:2] + (num >> 32) + (new_low < counter[0])
561 counter.assign(new_low.cat(new_high))
562 low = counter[0:1] - (num & 0xffffffff)
563 high = counter[1:2] - (num >> 32) - (counter[0] < (num & 0xffffffff))
564 return Tensor._device_seeds[device], low.cat(high)
565
566 @staticmethod
567 def rand(*shape, device:str|None=None, dtype:DTypeLike|None=None, contiguous:bool=True) -> Tensor:

Callers 1

randMethod · 0.80

Calls 5

TensorClass · 0.85
catMethod · 0.80
from_bytesMethod · 0.45
to_bytesMethod · 0.45
assignMethod · 0.45

Tested by

no test coverage detected