Files
tinygrad/extra/threefry.py
chenyu da10cf0be1 extra/threefry.py for mem usage (#4533)
for now it needs 8N mem to generate size N rand
2024-05-11 13:46:44 -04:00

10 lines
433 B
Python

import os
if "DEBUG" not in os.environ: os.environ["DEBUG"] = "2"
if "THREEFRY" not in os.environ: os.environ["THREEFRY"] = "1"
from tinygrad import Tensor, GlobalCounters
for N in [10, 100, 1_000, 10_000, 100_000, 1_000_000, 10_000_000, 100_000_000, 1_000_000_000]:
GlobalCounters.reset()
Tensor.rand(N).realize()
print(f"N {N:>20_}, global_ops {GlobalCounters.global_ops:>20_}, global_mem {GlobalCounters.global_mem:>20_}")