remove numpy from dtype (#4969)

replaced all dtype.np with _to_np_dtype defined in tensor.py.

after this, the only numpy usages are (1) Tensor(np.ndarray), (2) construct .numpy() output, (3) numpy random buffer

test/external/external_test_opt.py

@@ -7,6 +7,7 @@ from tinygrad import GlobalCounters, Tensor, Device
 from tinygrad.helpers import getenv
 from tinygrad.nn.state import get_parameters
 from tinygrad.engine.realize import capturing
+from tinygrad.tensor import _to_np_dtype
 
 PUSH_PERMUTES = False
 
@@ -46,21 +47,21 @@ class TestInferenceMinKernels(unittest.TestCase):
   @unittest.skipIf(not PUSH_PERMUTES, "this test requires PUSH_PERMUTES")
   def test_convnext(self):
     model = ConvNeXt()
-    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=p.dtype.np))
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
     img = Tensor.randn(1, 3, 224, 224)
     with CLCache(129):
       model(img).realize()
 
   def test_enet(self):
     model = EfficientNet(getenv("ENET_NUM", 0), has_se=False)
-    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=p.dtype.np))
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
     img = Tensor.randn(1, 3, 224, 224)
     with CLCache(51):
       model.forward(img).realize()
 
   def test_enet_se(self):
     model = EfficientNet(getenv("ENET_NUM", 0), has_se=True)
-    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=p.dtype.np))
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
     img = Tensor.randn(1, 3, 224, 224)
     # TODO: this seems very high
     with CLCache(115):
@@ -68,14 +69,14 @@ class TestInferenceMinKernels(unittest.TestCase):
 
   def test_resnet(self):
     model = ResNet18()
-    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=p.dtype.np))
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
     img = Tensor.randn(1, 3, 224, 224)
     with CLCache(23):
       model.forward(img).realize()
 
   def test_vit(self):
     model = ViT(embed_dim=192, num_heads=3)
-    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=p.dtype.np))
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
     img = Tensor.randn(1, 3, 224, 224)
     with CLCache(209) as cache: # NOTE: this is way too high
       out = model.forward(img)
@@ -87,7 +88,7 @@ class TestInferenceMinKernels(unittest.TestCase):
     from examples.llama import Transformer
     args_tiny = {"dim": 512, "hidden_dim": 1024, "n_heads": 8, "n_layers": 4, "norm_eps": 1e-05, "vocab_size": 1000}
     model = Transformer(**args_tiny)
-    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=p.dtype.np))
+    for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype)))
     inp = Tensor([[1,2,3,4]])
     with CLCache(100):
       model(inp, 0).realize()

test/external/fuzz_linearizer.py

@@ -5,6 +5,7 @@ from collections import defaultdict
 from extra.optimization.helpers import load_worlds, ast_str_to_lin
 
 from tinygrad import Tensor, Device, dtypes
+from tinygrad.tensor import _to_np_dtype
 from tinygrad.codegen.linearizer import Linearizer, UOp
 from tinygrad.codegen.kernel import Opt, OptOps
 from tinygrad.engine.search import get_linearizer_actions, bufs_from_lin
@@ -30,15 +31,15 @@ def get_fuzz_rawbufs(lin):
   with Context(DEBUG=0):
     for rawbuf in rawbufs[1:]:
       if dtypes.is_unsigned(rawbuf.dtype):
-        data = np.random.randint(0, 100, size=rawbuf.size, dtype=rawbuf.dtype.np)
+        data = np.random.randint(0, 100, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
       elif dtypes.is_int(rawbuf.dtype):
-        data = np.random.randint(-100, 100, size=rawbuf.size, dtype=rawbuf.dtype.np)
+        data = np.random.randint(-100, 100, size=rawbuf.size, dtype=_to_np_dtype(rawbuf.dtype))
       elif rawbuf.dtype == dtypes.bool:
         data = np.random.choice([True, False], size=rawbuf.size)
       elif rawbuf.dtype == dtypes.half:
-        data = np.random.uniform(-1, 1, size=rawbuf.size).astype(dtype=rawbuf.dtype.np)
+        data = np.random.uniform(-1, 1, size=rawbuf.size).astype(dtype=_to_np_dtype(rawbuf.dtype))
       else:
-        data = np.random.uniform(-10, 10, size=rawbuf.size).astype(dtype=rawbuf.dtype.np)
+        data = np.random.uniform(-10, 10, size=rawbuf.size).astype(dtype=_to_np_dtype(rawbuf.dtype))
       rawbuf.copyin(Tensor(data).realize().lazydata.realized.as_buffer())
   return rawbufs
 
@@ -92,7 +93,7 @@ def compare_linearizer(lin: Linearizer, rawbufs=None, var_vals=None, ground_trut
     unoptimized.required_optimizations()
     if run_linearizer(unoptimized, rawbufs, var_vals) != "PASS":
       return ("BASELINE_ERROR", rawbufs, var_vals, ground_truth,)
-    ground_truth = np.frombuffer(rawbufs[0].as_buffer(), rawbufs[0].dtype.np).copy()
+    ground_truth = np.frombuffer(rawbufs[0].as_buffer(), _to_np_dtype(rawbufs[0].dtype)).copy()
 
   rawbufs[0] = get_fuzz_rawbuf_like(rawbufs[0], zero=True) # get a new output buffer
   if (run_msg := run_linearizer(lin, rawbufs, var_vals)) != "PASS":
@@ -100,7 +101,7 @@ def compare_linearizer(lin: Linearizer, rawbufs=None, var_vals=None, ground_trut
 
   try:
     if not has_bf16:
-      result = np.frombuffer(rawbufs[0].as_buffer(), rawbufs[0].dtype.np)
+      result = np.frombuffer(rawbufs[0].as_buffer(), _to_np_dtype(rawbufs[0].dtype))
       np.testing.assert_allclose(result, ground_truth, rtol=rtol, atol=atol)
   except AssertionError as e:
     if DEBUG >= 2:

test/external/fuzz_schedule.py

@@ -7,7 +7,7 @@ from tinygrad.helpers import DEBUG, MULTIOUTPUT, colored, getenv
 from tinygrad.lazy import LazyBuffer
 from tinygrad.engine.schedule import _graph_schedule, _LBScheduleItem, ScheduleItem
 from tinygrad.ops import LoadOps
-from tinygrad.tensor import Tensor
+from tinygrad.tensor import Tensor, _to_np_dtype
 
 ctx_vars = { MULTIOUTPUT: (0, 1) }
 
@@ -60,8 +60,8 @@ def fuzz_schedule(outs:List[LazyBuffer]):
       si = ScheduleItem(ps.ast, tuple(rawbufs[x] for x in (ps.outputs+ps.inputs) if x.size != 0))
       _exec_si(si, seed)
       for out in ps.outputs:
-        outbuf = np.frombuffer(rawbufs[out].as_buffer(), out.dtype.np)
-        try: np.testing.assert_allclose(outbuf, np.frombuffer(ground_truth[out], out.dtype.np), atol=1e-2, rtol=1e-2)
+        outbuf = np.frombuffer(rawbufs[out].as_buffer(), _to_np_dtype(out.dtype))
+        try: np.testing.assert_allclose(outbuf, np.frombuffer(ground_truth[out], _to_np_dtype(out.dtype)), atol=1e-2, rtol=1e-2)
         except Exception as e:
           print(f"FAILED FOR {out}")
           raise e

test/external/fuzz_uops.py

@@ -6,6 +6,7 @@ from tinygrad.device import Buffer, Device
 from tinygrad.engine.realize import CompiledRunner
 from tinygrad.helpers import DEBUG, colored, getenv
 from tinygrad.shape.symbolic import Variable
+from tinygrad.tensor import _to_np_dtype
 
 def fuzz_uops(graph:DefaultDict[UOp, List[UOp]], in_degree:DefaultDict[UOp, int], loops_children:Dict[UOp, Set[UOp]]):
   paths: List[List[UOp]] = []
@@ -27,7 +28,7 @@ class UOpsFuzzerRunner(CompiledRunner):
     if DEBUG >= 1: print(colored(f"fuzzing {len(self.p.uops.fuzz_paths)} UOps permutations for {init_name}", "yellow"))
 
     super().__call__(rawbufs, var_vals, wait)
-    ground_truth = {x:np.frombuffer(x.as_buffer(), x.dtype.np) for x in rawbufs}
+    ground_truth = {x:np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)) for x in rawbufs}
 
     for i, path in enumerate(self.p.uops.fuzz_paths):
       # setup prg
@@ -43,7 +44,7 @@ class UOpsFuzzerRunner(CompiledRunner):
       super().__call__(rawbufs, var_vals, wait)
       for i, x in enumerate(rawbufs):
         try:
-          np.testing.assert_allclose(np.frombuffer(x.as_buffer(), x.dtype.np), ground_truth[x], atol=1e-6, rtol=1e-6)
+          np.testing.assert_allclose(np.frombuffer(x.as_buffer(), _to_np_dtype(x.dtype)), ground_truth[x], atol=1e-6, rtol=1e-6)
           if DEBUG >= 2: print(colored(name, "green"))
         except AssertionError as e:
           print(colored(name, "red"))

test/external/speed_compare_cuda_nv.py

@@ -4,6 +4,7 @@ from extra.optimization.helpers import load_worlds, ast_str_to_lin
 from test.external.fuzz_linearizer import get_fuzz_rawbufs
 from tinygrad.engine.search import bufs_from_lin
 from tinygrad.engine.realize import CompiledRunner
+from tinygrad.tensor import _to_np_dtype
 import numpy as np
 
 # move to helpers?
@@ -63,8 +64,8 @@ if __name__ == "__main__":
       failed = True
 
     if not failed and not has_bf16:
-      curesult = np.frombuffer(test_cubufs[0].as_buffer(), test_cubufs[0].dtype.np)
-      nvresult = np.frombuffer(test_nvbufs[0].as_buffer(), test_nvbufs[0].dtype.np)
+      curesult = np.frombuffer(test_cubufs[0].as_buffer(), _to_np_dtype(test_cubufs[0].dtype))
+      nvresult = np.frombuffer(test_nvbufs[0].as_buffer(), _to_np_dtype(test_nvbufs[0].dtype))
       np.testing.assert_allclose(curesult, nvresult, rtol=1e-2, atol=1e-2)
 
     average_tm_cuda += min(tm_cuda)