github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-10 07:28:15 -05:00
Code Issues Packages Projects Releases Wiki Activity

simple_matmul.py uses np to generate random (#9438)

Browse Source 
* np generates randoms

* hotfix: use generator for int dtype

* float32 as default dtype for float generator

* use np.float32 instead of stirng

* add dtype= to integers generator

* change import _to_np_dtype source
This commit is contained in:
Ignacio Sica
2025-03-14 17:36:50 -03:00
committed by GitHub
parent 2a50e6440d
commit 14018050c1
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
extra/gemm/simple_matmul.py
View File

@@ -1,5 +1,6 @@
import numpy as np
from tinygrad.helpers import getenv
from tinygrad.dtype import _to_np_dtype
from tinygrad import dtypes, Tensor
dtype_in = dtypes.half if getenv("HALF") else dtypes.bfloat16 if getenv("BFLOAT16") else dtypes.float
@@ -13,12 +14,15 @@ K = getenv("K", N)
CNT = getenv("CNT", 10)
ATOL = getenv("ATOL", 1e-4)
RTOL = getenv("RTOL", 3e-2)
INT_LOW = getenv("INT_LOW", 0)
INT_HIGH = getenv("INT_HIGH", 10)
if __name__ == "__main__":
def init_matrix(rows, cols):
rng = np.random.default_rng()
if dtype_in in dtypes.ints:
return Tensor.randint((rows, cols), dtype=dtype_in).realize()
return Tensor.rand(rows, cols, dtype=dtype_in).realize()
return Tensor(rng.integers(INT_LOW, INT_HIGH, (rows, cols), dtype=_to_np_dtype(dtype_in))).realize()
return Tensor(rng.random((rows, cols), dtype=np.float32).astype(_to_np_dtype(dtype_in))).realize()
a, b = init_matrix(M, K), init_matrix(K, N)
for i in range(CNT):