tinygrad 0.9.1__tar.gz → 0.10.0__tar.gz

{tinygrad-0.9.1 → tinygrad-0.10.0}/PKG-INFO

@@ -1,17 +1,14 @@
 Metadata-Version: 2.1
 Name: tinygrad
-Version: 0.9.1
+Version: 0.10.0
 Summary: You like pytorch? You like micrograd? You love tinygrad! <3
 Author: George Hotz
 License: MIT
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
-Requires-Python: >=3.8
+Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: numpy
-Requires-Dist: pyobjc-framework-Metal; platform_system == "Darwin"
-Requires-Dist: pyobjc-framework-libdispatch; platform_system == "Darwin"
 Provides-Extra: llvm
 Requires-Dist: llvmlite; extra == "llvm"
 Provides-Extra: arm
@@ -20,12 +17,13 @@ Provides-Extra: triton
 Requires-Dist: triton-nightly>=2.1.0.dev20231014192330; extra == "triton"
 Provides-Extra: linting
 Requires-Dist: pylint; extra == "linting"
-Requires-Dist: mypy; extra == "linting"
+Requires-Dist: mypy==1.11.2; extra == "linting"
 Requires-Dist: typing-extensions; extra == "linting"
 Requires-Dist: pre-commit; extra == "linting"
 Requires-Dist: ruff; extra == "linting"
 Requires-Dist: types-tqdm; extra == "linting"
 Provides-Extra: testing
+Requires-Dist: numpy; extra == "testing"
 Requires-Dist: torch; extra == "testing"
 Requires-Dist: pillow; extra == "testing"
 Requires-Dist: pytest; extra == "testing"
@@ -39,17 +37,21 @@ Requires-Dist: safetensors; extra == "testing"
 Requires-Dist: transformers; extra == "testing"
 Requires-Dist: sentencepiece; extra == "testing"
 Requires-Dist: tiktoken; extra == "testing"
+Requires-Dist: blobfile; extra == "testing"
 Requires-Dist: librosa; extra == "testing"
 Requires-Dist: networkx; extra == "testing"
 Requires-Dist: hypothesis; extra == "testing"
 Requires-Dist: nibabel; extra == "testing"
 Requires-Dist: bottle; extra == "testing"
+Requires-Dist: ggml-python; extra == "testing"
 Provides-Extra: docs
+Requires-Dist: mkdocs; extra == "docs"
 Requires-Dist: mkdocs-material; extra == "docs"
 Requires-Dist: mkdocstrings[python]; extra == "docs"
 Requires-Dist: markdown-callouts; extra == "docs"
 Requires-Dist: markdown-exec[ansi]; extra == "docs"
 Requires-Dist: black; extra == "docs"
+Requires-Dist: numpy; extra == "docs"
 Provides-Extra: testing-tf
 Requires-Dist: tensorflow==2.15.1; extra == "testing-tf"
 Requires-Dist: tensorflow_addons; extra == "testing-tf"
@@ -107,7 +109,7 @@ And we can change `DEBUG` to `4` to see the generated code.
 As it turns out, 90% of what you need for neural networks are a decent autograd/tensor library.
 Throw in an optimizer, a data loader, and some compute, and you have all you need.
 
-```py
+```python
 from tinygrad import Tensor, nn
 
 class LinearNet:
@@ -122,11 +124,12 @@ optim = nn.optim.Adam([model.l1, model.l2], lr=0.001)
 
 x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7])  # replace with real mnist dataloader
 
-for i in range(10):
-  optim.zero_grad()
-  loss = model(x).sparse_categorical_crossentropy(y).backward()
-  optim.step()
-  print(i, loss.item())
+with Tensor.train():
+  for i in range(10):
+    optim.zero_grad()
+    loss = model(x).sparse_categorical_crossentropy(y).backward()
+    optim.step()
+    print(i, loss.item())
 ```
 
 See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full version that gets 98% in ~5 seconds
@@ -142,9 +145,12 @@ tinygrad already supports numerous accelerators, including:
 - [x] [CUDA](tinygrad/runtime/ops_cuda.py)
 - [x] [AMD](tinygrad/runtime/ops_amd.py)
 - [x] [NV](tinygrad/runtime/ops_nv.py)
+- [x] [QCOM](tinygrad/runtime/ops_qcom.py)
 
 And it is easy to add more! Your accelerator of choice only needs to support a total of ~25 low level ops.
 
+To check default accelerator run: `python3 -c "from tinygrad import Device; print(Device.DEFAULT)"`
+
 ## Installation
 
 The current recommended way to install tinygrad is from source.
@@ -169,7 +175,7 @@ Documentation along with a quick start guide can be found on the [docs website](
 
 ### Quick example comparing to PyTorch
 
-```py
+```python
 from tinygrad import Tensor
 
 x = Tensor.eye(3, requires_grad=True)
@@ -182,7 +188,7 @@ print(y.grad.numpy())  # dz/dy
 ```
 
 The same thing but in PyTorch:
-```py
+```python
 import torch
 
 x = torch.eye(3, requires_grad=True)
@@ -230,6 +236,4 @@ python3 -m pytest test/                 # whole test suite
 
 #### Process replay tests
 
-[Process replay](https://github.com/tinygrad/tinygrad/blob/master/test/external/process_replay/process_replay.py) detects changes in the generated kernels of CI tests by comparing them against tinygrad master. If your PR is a refactor or speedup without any expected behavior change, it should include a green process replay pass to get merged.
-
-You can enable process replay by adding [run_process_replay] to your PR title. [example](https://github.com/tinygrad/tinygrad/pull/4995). Note that you should keep your branch up-to-date with master.
+[Process replay](https://github.com/tinygrad/tinygrad/blob/master/test/external/process_replay/README.md) compares your PR's generated kernels against master. If your PR is a refactor or speedup without any expected behavior change, It should include [pr] in the pull request title.

{tinygrad-0.9.1 → tinygrad-0.10.0}/README.md

@@ -51,7 +51,7 @@ And we can change `DEBUG` to `4` to see the generated code.
 As it turns out, 90% of what you need for neural networks are a decent autograd/tensor library.
 Throw in an optimizer, a data loader, and some compute, and you have all you need.
 
-```py
+```python
 from tinygrad import Tensor, nn
 
 class LinearNet:
@@ -66,11 +66,12 @@ optim = nn.optim.Adam([model.l1, model.l2], lr=0.001)
 
 x, y = Tensor.rand(4, 1, 28, 28), Tensor([2,4,3,7])  # replace with real mnist dataloader
 
-for i in range(10):
-  optim.zero_grad()
-  loss = model(x).sparse_categorical_crossentropy(y).backward()
-  optim.step()
-  print(i, loss.item())
+with Tensor.train():
+  for i in range(10):
+    optim.zero_grad()
+    loss = model(x).sparse_categorical_crossentropy(y).backward()
+    optim.step()
+    print(i, loss.item())
 ```
 
 See [examples/beautiful_mnist.py](examples/beautiful_mnist.py) for the full version that gets 98% in ~5 seconds
@@ -86,9 +87,12 @@ tinygrad already supports numerous accelerators, including:
 - [x] [CUDA](tinygrad/runtime/ops_cuda.py)
 - [x] [AMD](tinygrad/runtime/ops_amd.py)
 - [x] [NV](tinygrad/runtime/ops_nv.py)
+- [x] [QCOM](tinygrad/runtime/ops_qcom.py)
 
 And it is easy to add more! Your accelerator of choice only needs to support a total of ~25 low level ops.
 
+To check default accelerator run: `python3 -c "from tinygrad import Device; print(Device.DEFAULT)"`
+
 ## Installation
 
 The current recommended way to install tinygrad is from source.
@@ -113,7 +117,7 @@ Documentation along with a quick start guide can be found on the [docs website](
 
 ### Quick example comparing to PyTorch
 
-```py
+```python
 from tinygrad import Tensor
 
 x = Tensor.eye(3, requires_grad=True)
@@ -126,7 +130,7 @@ print(y.grad.numpy())  # dz/dy
 ```
 
 The same thing but in PyTorch:
-```py
+```python
 import torch
 
 x = torch.eye(3, requires_grad=True)
@@ -174,6 +178,4 @@ python3 -m pytest test/                 # whole test suite
 
 #### Process replay tests
 
-[Process replay](https://github.com/tinygrad/tinygrad/blob/master/test/external/process_replay/process_replay.py) detects changes in the generated kernels of CI tests by comparing them against tinygrad master. If your PR is a refactor or speedup without any expected behavior change, it should include a green process replay pass to get merged.
-
-You can enable process replay by adding [run_process_replay] to your PR title. [example](https://github.com/tinygrad/tinygrad/pull/4995). Note that you should keep your branch up-to-date with master.
+[Process replay](https://github.com/tinygrad/tinygrad/blob/master/test/external/process_replay/README.md) compares your PR's generated kernels against master. If your PR is a refactor or speedup without any expected behavior change, It should include [pr] in the pull request title.

{tinygrad-0.9.1 → tinygrad-0.10.0}/setup.py

@@ -8,36 +8,36 @@ with open(directory / 'README.md', encoding='utf-8') as f:
   long_description = f.read()
 
 setup(name='tinygrad',
-      version='0.9.1',
+      version='0.10.0',
       description='You like pytorch? You like micrograd? You love tinygrad! <3',
       author='George Hotz',
       license='MIT',
       long_description=long_description,
       long_description_content_type='text/markdown',
       packages = ['tinygrad', 'tinygrad.runtime.autogen', 'tinygrad.codegen', 'tinygrad.nn', 'tinygrad.renderer', 'tinygrad.engine',
-                  'tinygrad.runtime', 'tinygrad.runtime.driver', 'tinygrad.runtime.graph', 'tinygrad.shape'],
+                  'tinygrad.runtime', 'tinygrad.runtime.support', 'tinygrad.runtime.graph', 'tinygrad.shape'],
+      package_data = {'tinygrad': ['py.typed']},
       classifiers=[
         "Programming Language :: Python :: 3",
         "License :: OSI Approved :: MIT License"
       ],
-      install_requires=["numpy",
-                        "pyobjc-framework-Metal; platform_system=='Darwin'",
-                        "pyobjc-framework-libdispatch; platform_system=='Darwin'"],
-      python_requires='>=3.8',
+      install_requires=[],
+      python_requires='>=3.10',
       extras_require={
         'llvm': ["llvmlite"],
         'arm': ["unicorn"],
         'triton': ["triton-nightly>=2.1.0.dev20231014192330"],
         'linting': [
             "pylint",
-            "mypy",
+            "mypy==1.11.2",
             "typing-extensions",
             "pre-commit",
             "ruff",
             "types-tqdm",
         ],
-        #'mlperf': ["mlperf-logging @ git+https://github.com/mlperf/logging.git@4.0.0-rc2"],
+        #'mlperf': ["mlperf-logging @ git+https://github.com/mlperf/logging.git@4.1.0-rc3"],
         'testing': [
+            "numpy",
             "torch",
             "pillow",
             "pytest",
@@ -51,18 +51,22 @@ setup(name='tinygrad',
             "transformers",
             "sentencepiece",
             "tiktoken",
+            "blobfile",
             "librosa",
             "networkx",
             "hypothesis",
             "nibabel",
             "bottle",
+            "ggml-python"
         ],
         'docs': [
+            "mkdocs",
             "mkdocs-material",
             "mkdocstrings[python]",
             "markdown-callouts",
             "markdown-exec[ansi]",
-            "black"
+            "black",
+            "numpy",
         ],
         'testing_tf': [
             "tensorflow==2.15.1",

tinygrad-0.10.0/test/test_arange.py

@@ -0,0 +1,179 @@
+import unittest, contextlib
+import numpy as np
+from tinygrad import Tensor, GlobalCounters, dtypes, nn
+from tinygrad.helpers import CI, Context, getenv
+from tinygrad.engine.realize import run_schedule
+from tinygrad.codegen.kernel import Opt, OptOps, Kernel, KernelOptError
+from tinygrad.engine.realize import CompiledRunner, ExecItem
+from tinygrad.engine.search import get_kernel_actions
+
+class TestArange(unittest.TestCase):
+  def _get_flops(self, N, opts=None):
+    GlobalCounters.reset()
+    tt = Tensor.arange(N)
+    sched = tt.schedule()
+    self.assertEqual(len(sched), 1)
+    k = Kernel(sched[-1].ast)
+    if opts is not None:
+      for o in opts: k.apply_opt(o)
+    p = k.to_program()
+    print(p.name)
+    #print(p.src)
+    ExecItem(CompiledRunner(p), [tt.lazydata.buffer]).run()
+    np.testing.assert_equal(tt.numpy(), np.arange(N))
+    return p.op_estimate
+
+  def test_complexity(self, opts=None, limit=None):
+    # add 1 to avoid divide by 0. arange is 0 flops now!
+    f1 = self._get_flops(256, opts) + 1
+    f2 = self._get_flops(2560, opts) + 1
+    print(f"{f1=}, {f2=}")
+    assert (f1 < 5000 and f2 < 5000) or (f2 / f1 < 15), f"bad complexity, flops {f2/f1:.1f}X while inputs 10X"
+    if limit is not None and not getenv("PTX"):
+      # PTX counts index ALU in flops
+      assert f1 <= limit, f"{f1=}, {limit=}"
+
+  def test_complexity_w_upcast(self): return self.test_complexity([Opt(OptOps.UPCAST, 0, 4)], limit=1)
+  def test_complexity_w_unroll2(self): return self.test_complexity([Opt(OptOps.UNROLL, 0, 2)], limit=1)
+  def test_complexity_w_unroll4(self): return self.test_complexity([Opt(OptOps.UNROLL, 0, 4)], limit=1)
+  def test_complexity_w_unroll8(self): return self.test_complexity([Opt(OptOps.UNROLL, 0, 8)], limit=1)
+  def test_complexity_w_upcast_and_unroll(self): return self.test_complexity([Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)], limit=1)
+
+  @unittest.skip("doesn't work yet")
+  def test_complexity_w_local_and_padto(self): return self.test_complexity([Opt(OptOps.LOCAL, 0, 16), Opt(op=OptOps.PADTO, axis=1, amt=32)])
+
+  def test_all_opts(self, opts=None, exclude=None):
+    k = Kernel(Tensor.arange(256).schedule()[-1].ast)
+    if opts is not None:
+      for o in opts: k.apply_opt(o)
+    all_opts_256 = [kk.applied_opts for kk in get_kernel_actions(k, include_0=False).values()]
+    k = Kernel(Tensor.arange(2560).schedule()[-1].ast)
+    if opts is not None:
+      for o in opts: k.apply_opt(o)
+    all_opts_2560 = [kk.applied_opts for kk in get_kernel_actions(k, include_0=False).values()]
+    all_opts = [x for x in all_opts_256 if x in all_opts_2560]
+    for opts in all_opts:
+      if exclude is not None and opts[-1] in exclude: continue
+      print(opts)
+      self.test_complexity(opts)
+  def test_all_opts_w_local(self):
+    with contextlib.suppress(KernelOptError):
+      return self.test_all_opts([Opt(OptOps.LOCAL, 0, 16)], [Opt(op=OptOps.PADTO, axis=1, amt=32)])
+  def test_all_opts_w_upcast(self): return self.test_all_opts([Opt(OptOps.UPCAST, 0, 4)])
+  def test_all_opts_w_unroll(self): return self.test_all_opts([Opt(OptOps.UNROLL, 0, 4)], [Opt(op=OptOps.GROUP, axis=0, amt=0)])
+  def test_all_opts_w_upcast_and_unroll(self):
+    return self.test_all_opts([Opt(OptOps.UPCAST, 0, 4), Opt(OptOps.UNROLL, 0, 4)], [Opt(op=OptOps.GROUP, axis=0, amt=0)])
+
+class TestIndexing(unittest.TestCase):
+  def test_arange_2_reduce(self):
+    needle = Tensor.zeros(16384, dtype=dtypes.int).contiguous()
+    needle[1337] = 1
+    needle.realize()
+    with Context(NOOPT=1, FUSE_ARANGE=1):
+      GlobalCounters.reset()
+      # TODO: it should work without these reshapes
+      out = ((Tensor.arange(1,16385).reshape(16384,1)-1)*needle.reshape(16384,1)).sum()
+      sched = out.schedule()
+      assert len(sched) == 1
+      run_schedule(sched)
+    assert out.item() == 1337, f"expected 1337, got {out.item()}"
+
+  @unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")
+  def test_manual_index(self):
+    dataset = Tensor.rand(16384, 256).realize()
+    idxs = Tensor([0,3,5,6]).realize()
+    real_index = dataset.numpy()[idxs.numpy()]
+    print("*** indexing ***")
+    with Context(NOOPT=1, FUSE_ARANGE=1):
+      GlobalCounters.reset()
+      rng = Tensor.ones(4, 256, 16384, dtype=dtypes.int)._cumsum(axis=-1, _first_zero=True).reshape(4, 256, 16384, 1)
+      idxs = idxs.reshape(4,1,1,1).expand(4, 256, 16384, 1)
+      reshape_dataset = dataset.T.reshape(1, 256, 16384, 1).expand(4, 256, 16384, 1)
+      full = (rng==idxs).where(reshape_dataset, Tensor.zeros(4, 256, 16384, 1))
+      X = full.sum(axis=(2,3))
+      sched = X.schedule()
+      assert len(sched) == 1
+      run_schedule(sched)
+      assert GlobalCounters.global_ops < 4*16384, f"too many ops {GlobalCounters.global_ops}"
+    np.testing.assert_allclose(real_index, X.numpy())
+
+  def test_index(self):
+    dataset = Tensor.rand(16384, 256).realize()
+    idxs = Tensor([0,3,5,6]).realize()
+    real_index = dataset.numpy()[idxs.numpy()]
+    print("*** indexing ***")
+    with Context(NOOPT=1):
+      GlobalCounters.reset()
+      X = dataset[idxs]
+      assert X.shape == (4,256)
+      sched = X.schedule()
+      # TODO: enable these asserts when the scheduler can handle this
+      #assert len(sched) == 1, f"{len(sched)} != 1"
+      run_schedule(sched)
+      #assert GlobalCounters.global_ops < 4*16384, f"too many ops {GlobalCounters.global_ops}"
+    np.testing.assert_allclose(real_index, X.numpy())
+
+  def test_index_fused(self, noopt=1):
+    dataset = Tensor.rand(16384, 256).realize()
+    idxs = Tensor([0,3,5,6]).realize()
+    real_index = dataset.numpy()[idxs.numpy()]
+    print("*** indexing ***")
+    with Context(NOOPT=noopt, FUSE_ARANGE=1):
+      GlobalCounters.reset()
+      X = dataset[idxs]
+      assert X.shape == (4,256)
+      sched = X.schedule()
+      assert len(sched) == 2
+      run_schedule(sched)
+      assert GlobalCounters.global_ops < 4*16384, f"too many ops {GlobalCounters.global_ops} != {4*16384}"
+    np.testing.assert_allclose(real_index, X.numpy())
+  @unittest.skip("not ready")
+  def test_index_fused_opt(self): self.test_index_fused(0)
+
+  def test_index_fused_out_of_bounds(self):
+    dataset = Tensor.rand(256, 256).realize()
+    idxs = Tensor([-19238, -257, 256, 495, 10982377]).realize()
+    with Context(NOOPT=1, FUSE_ARANGE=1):
+      X = dataset[idxs]
+      np.testing.assert_equal(X.numpy(), 0)
+
+  @unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")
+  def test_index_mnist(self, noopt=1, op_limit=512*784*5):
+    from tinygrad.nn.datasets import mnist
+    X_train, Y_train, _, _ = mnist()
+    with Context(NOOPT=noopt, FUSE_ARANGE=1, SPLIT_REDUCEOP=0):
+      samples = Tensor.randint(getenv("BS", 512), high=X_train.shape[0]).realize()
+      GlobalCounters.reset()
+      x = X_train[samples].numpy()
+      y = Y_train[samples].numpy()
+      assert GlobalCounters.global_ops < op_limit, f"too many ops {GlobalCounters.global_ops} != {op_limit}"
+    np.testing.assert_allclose(X_train.numpy()[samples.numpy()], x)
+    np.testing.assert_allclose(Y_train.numpy()[samples.numpy()], y)
+  @unittest.skip("not ready")
+  def test_index_mnist_opt(self): self.test_index_mnist(0)
+
+  @unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")
+  def test_llama_embedding(self, noopt=1, op_limit=65536):
+    # llama3 is 128256
+    vocab_size, embed_size = (10, 3) if CI else (32000, 4096)
+    emb = nn.Embedding(vocab_size, embed_size)
+    emb_w = emb.weight.numpy()
+    x = Tensor([1,2,3,4])
+    with Context(NOOPT=noopt, FUSE_ARANGE=1):
+      GlobalCounters.reset()
+      z = emb(x).realize()
+      self.assertLessEqual(GlobalCounters.global_ops, op_limit)
+      self.assertEqual(GlobalCounters.kernel_count, 2)
+    if getenv("CHECK", 1):
+      import torch
+      with torch.no_grad():
+        torch_emb = torch.nn.Embedding(vocab_size, embed_size).eval()
+        torch_emb.weight[:] = torch.tensor(emb_w, dtype=torch.float32)
+      torch_z = torch_emb(torch.tensor(x.numpy()))
+      # TODO: reshape to match torch, should we do this in nn?
+      np.testing.assert_allclose(z.numpy().reshape(4, embed_size), torch_z.detach().numpy(), atol=1e-8, rtol=1e-8)
+  # at least the arange is being fused
+  def test_llama_embedding_opt(self): self.test_llama_embedding(0, 1_736_704_000 if CI else 5_898_240_000)
+
+if __name__ == "__main__":
+  unittest.main()

{tinygrad-0.9.1 → tinygrad-0.10.0}/test/test_assign.py

@@ -2,6 +2,7 @@
 import unittest
 import numpy as np
 from tinygrad import dtypes, Tensor, TinyJit, GlobalCounters, Variable
+from tinygrad.engine.schedule import create_schedule
 
 N = 200  # has to be bigger than the cache to fail
 
@@ -57,10 +58,12 @@ class TestAssign(unittest.TestCase):
       x.realize()
     x = Tensor([0])
     f(x)
-    assert (out:=x.item()) == 1, f"expected 1, got {out}"
+    out = x.item()
+    assert out == 1, f"expected 1, got {out}"
     x = Tensor([0])
     f(x)
-    assert (out:=x.item()) == 1, f"expected 1, got {out}"
+    out = x.item()
+    assert out == 1, f"expected 1, got {out}"
 
   def test_assign_add_jit(self):
     @TinyJit
@@ -165,6 +168,16 @@ class TestAssign(unittest.TestCase):
     a += 1
     np.testing.assert_allclose(a.numpy(), 3)
 
+  # NOTE: this is similar to the resnet failure
+  #@unittest.expectedFailure
+  def test_double_assign_alt(self):
+    a = Tensor.ones(4).contiguous().realize()
+    b = Tensor([1, 2, 3, 4]).realize().lazydata
+    a1 = a.lazydata.assign(b)
+    a2 = a.lazydata.assign(b)
+    sched = create_schedule([a1, a2])
+    self.assertEqual(len(sched), 1)
+
   def test_crossover_assign(self):
     a = Tensor.full((4,), 2).contiguous().realize()
     b = Tensor.full((4,), 3).contiguous().realize()
@@ -347,7 +360,7 @@ class TestAssign(unittest.TestCase):
 
   def test_permuted_assignment_masked_view_possible(self):
     a = Tensor.ones(4, 4).contiguous().realize()
-    b = a.shrink((None, (0, 2))).pad((None, (0, 2)), 2)
+    b = a.shrink((None, (0, 2))).pad((None, (0, 2)), value=2)
     a.assign(a + b)
     kc = GlobalCounters.kernel_count
     a.realize()
@@ -357,7 +370,7 @@ class TestAssign(unittest.TestCase):
   def test_permuted_assignment_masked_view_not_contiguous(self):
     a = Tensor.ones(4, 4).contiguous().realize()
     with self.assertRaisesRegex(RuntimeError, "contiguous"):
-      b = a.shrink((None, (0, 2))).pad((None, (0, 2)), 2).permute(1, 0)
+      b = a.shrink((None, (0, 2))).pad((None, (0, 2)), value=2).permute(1, 0)
       a.assign(a + b)
       a.realize()
 

tinygrad-0.10.0/test/test_compile_failures.py

@@ -0,0 +1,18 @@
+import unittest
+from tinygrad import Tensor, dtypes, Device
+from tinygrad.engine.realize import lower_schedule
+from tinygrad.device import is_dtype_supported
+
+class TestCompileFailures(unittest.TestCase):
+  def compile(self, out:Tensor):
+    for _ in lower_schedule(out.schedule()): pass
+
+  @unittest.skipUnless(is_dtype_supported(dtypes.uchar, Device.DEFAULT), f"no uint8 on {Device.DEFAULT}")
+  def test_interpolate_atari(self):
+    self.compile(Tensor.empty(210, 160, dtype='uint8').interpolate((64, 64)))
+
+  def test_add_max_uchar(self):
+    self.compile((Tensor.empty(1024, dtype='uint8') + Tensor.empty(1024, dtype='uint8')).max())
+
+if __name__ == '__main__':
+  unittest.main()

{tinygrad-0.9.1 → tinygrad-0.10.0}/test/test_const_folding.py

@@ -1,15 +1,15 @@
 import unittest, math
 from tinygrad import Tensor, Device, dtypes
+from tinygrad.ops import Ops
 from tinygrad.engine.schedule import create_schedule
 from tinygrad.helpers import CI
-from tinygrad.ops import BufferOps
 import numpy as np
-from test.helpers import is_dtype_supported
+from tinygrad.device import is_dtype_supported
 
 def _check_ast_count(desired_count:int, t:Tensor):
   # NOTE: this has side effect because everything can be scheduled only once
   schedule = create_schedule(t.lazydata.lbs)
-  asts = [s for s in schedule if s.ast[0].op is BufferOps.STORE]
+  asts = [s for s in schedule if s.ast.op is Ops.SINK]
   assert len(asts) == desired_count
 
 class TestUnaryOpsConstFolding(unittest.TestCase):
@@ -23,6 +23,7 @@ class TestUnaryOpsConstFolding(unittest.TestCase):
     _check_ast_count(0, Tensor.ones(4).cast(dtypes.int16))
     _check_ast_count(0, Tensor.full(4, fill_value=-1).cast(dtypes.uint16))
 
+  @unittest.expectedFailure  # no two level fold at lazybuffer
   def test_neg_folding(self):
     _check_ast_count(0, Tensor([1, 2, 3]).mul(-1).neg())
     _check_ast_count(0, Tensor([1, 2, 3]).neg().mul(-1))
@@ -78,6 +79,11 @@ class TestBinaryOpsConstFolding(unittest.TestCase):
   def test_div_tensor_one(self):
     _check_ast_count(0, Tensor([1.0, 2, 3, 4]) / Tensor.ones(4))
 
+  def test_idiv_literal_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) // 1)
+  def test_idiv_tensor_one(self):
+    _check_ast_count(0, Tensor([1, 2, 3, 4]) // Tensor.ones(4, dtype=dtypes.int32))
+
   def test_pow_literal_zero(self):
     _check_ast_count(0, Tensor([1.0, 2, 3, 4]) ** 0)
   def test_pow_tensor_zero(self):
@@ -124,13 +130,16 @@ class TestMovedConstFolding(unittest.TestCase):
 
   def test_cast_padded(self):
     # NOTE: this is folded due to CAST_BEFORE_VIEW
-    _check_ast_count(0, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
-    np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0])
-    _check_ast_count(0, Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16))
-    np.testing.assert_equal(Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16).numpy(), [0, 65535, 65535, 65535, 65535, 0])
+    if is_dtype_supported(dtypes.int16):
+      _check_ast_count(0, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16))
+      np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int16).numpy(), [0, 1, 1, 1, 1, 0])
+    if is_dtype_supported(dtypes.uint16):
+      _check_ast_count(0, Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16))
+      np.testing.assert_equal(Tensor.full(4, fill_value=-1).pad(((1, 1),)).cast(dtypes.uint16).numpy(), [0, 65535, 65535, 65535, 65535, 0])
     # not folded
-    _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64))
-    np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64).numpy(), [0, 1, 1, 1, 1, 0])
+    if is_dtype_supported(dtypes.int64):
+      _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64))
+      np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).cast(dtypes.int64).numpy(), [0, 1, 1, 1, 1, 0])
 
 class TestReduceOpsConstFolding(unittest.TestCase):
   def test_const_sum(self):
@@ -145,10 +154,18 @@ class TestReduceOpsConstFolding(unittest.TestCase):
     _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).sum())
     np.testing.assert_equal(Tensor.ones(4).pad(((1, 1),)).sum().numpy(), 4)
 
-    # NOTE: cannot just count the non-padded area because some UnaryOps f do not have f(0) = 0.
+    # NOTE: cannot just count the non-padded area because some Ops f do not have f(0) = 0.
     _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).exp().sum())
     np.testing.assert_allclose(Tensor.ones(4).pad(((1, 1),)).exp().sum().numpy(), 4 * math.e + 2)
 
+  def test_const_prod(self):
+    _check_ast_count(0, Tensor.full((2, 3), fill_value=2).prod())
+    np.testing.assert_equal(Tensor.full((2, 3), fill_value=2).prod().numpy(), 2**(2*3))
+    _check_ast_count(0, Tensor.full((4, 5, 6), fill_value=2).prod(axis=0))
+    np.testing.assert_equal(Tensor.full((4, 5, 6), fill_value=2).prod(axis=0).numpy(), np.full((5, 6), 2**4))
+    _check_ast_count(0, Tensor(4).prod())
+    np.testing.assert_equal(Tensor(4).prod().numpy(), 4)
+
   def test_const_max(self):
     _check_ast_count(0, Tensor.ones(4, 5, 6).max())
     np.testing.assert_equal(Tensor.ones(4, 5, 6).max().numpy(), 1)
@@ -234,7 +251,6 @@ class TestTautologicalCompare(unittest.TestCase):
     np.testing.assert_equal((Tensor(True) < Tensor(False)).numpy(), False)
     np.testing.assert_equal((Tensor(True) < Tensor(True)).numpy(), False)
 
-  @unittest.skip("not implemented yet")
   def test_a_eq_a(self):
     # self eq is always true for int or bool
     a = Tensor([1, 2, 3])
@@ -244,7 +260,6 @@ class TestTautologicalCompare(unittest.TestCase):
     a = Tensor([math.nan, 1.0, 2.0])
     np.testing.assert_equal((a == a).numpy(), [False, True, True])
 
-  @unittest.skip("not implemented yet")
   def test_a_ne_a(self):
     # self not eq is always false for int or bool
     a = Tensor([1, 2, 3])

{tinygrad-0.9.1 → tinygrad-0.10.0}/test/test_conv.py

@@ -42,6 +42,14 @@ class TestConv(unittest.TestCase):
 
     print(ret.numpy())
 
+  def test_two_binops_no_rerun_small(self):
+    Tensor.no_grad = True
+    x = Tensor.rand(1,1,32,32)
+    w = Tensor.rand(1,1,3,3)
+    out = x.conv2d(w, padding=(1,1))
+    np.testing.assert_allclose(out.relu().numpy(), np.maximum(out.numpy(), 0))
+    Tensor.no_grad = False
+
   def test_two_binops_no_rerun(self):
     Tensor.no_grad = True
     x = Tensor.randn(1,12,128,256)