tests in pre commit

.pre-commit-config.yaml

@@ -1,12 +1,6 @@
 repos:
   - repo: local
     hooks:
-      - id: pylint
-        name: pylint
-        entry: pylint tinygrad/
-        language: system
-        always_run: true
-        pass_filenames: false
       - id: flake8
         name: flake8
         entry: flake8 tinygrad/ --indent-size=2 --select=F,E112,E113,E203,E304,E502,E702,E703,E71,E72,E731,W191,W6 --statistics -j4
@@ -19,3 +13,15 @@ repos:
         language: system
         always_run: true
         pass_filenames: false
+      - id: tests
+        name: subset of tests
+        entry: pytest test/unit/ test/test_ops.py
+        language: system
+        always_run: true
+        pass_filenames: false
+      - id: pylint
+        name: pylint
+        entry: pylint tinygrad/
+        language: system
+        always_run: true
+        pass_filenames: false

test/test_example.py

@@ -1,24 +0,0 @@
-import unittest
-from tinygrad.tensor import Tensor
-
-class TestExample(unittest.TestCase):
-  def test_example_readme(self):
-    x = Tensor.eye(3, requires_grad=True)
-    y = Tensor([[2.0,0,-2.0]], requires_grad=True)
-    z = y.matmul(x).sum()
-    z.backward()
-
-    print(x.grad.numpy())  # dz/dx
-    print(y.grad.numpy())  # dz/dy
-
-  def test_example_matmul(self):
-    x = Tensor.eye(256, requires_grad=True)
-    y = Tensor.eye(256, requires_grad=True)
-    z = y.matmul(x).sum()
-    z.backward()
-
-    print(x.grad.numpy())  # dz/dx
-    print(y.grad.numpy())  # dz/dy
-
-if __name__ == '__main__':
-  unittest.main()

test/unit/test_example.py

@@ -0,0 +1,40 @@
+import unittest
+from tinygrad.tensor import Tensor
+
+class TestExample(unittest.TestCase):
+  def _test_example_readme(self, device):
+    x = Tensor.eye(3, device=device, requires_grad=True)
+    y = Tensor([[2.0,0,-2.0]], device=device, requires_grad=True)
+    z = y.matmul(x).sum()
+    z.backward()
+
+    print(x.grad.numpy())  # dz/dx
+    print(y.grad.numpy())  # dz/dy
+
+    assert x.grad.device == device
+    assert y.grad.device == device
+
+  def _test_example_matmul(self, device):
+    x = Tensor.eye(64, device=device, requires_grad=True)
+    y = Tensor.eye(64, device=device, requires_grad=True)
+    z = y.matmul(x).sum()
+    z.backward()
+
+    print(x.grad.numpy())  # dz/dx
+    print(y.grad.numpy())  # dz/dy
+
+    assert x.grad.device == device
+    assert y.grad.device == device
+
+  def test_example_readme_cpu(self): self._test_example_readme("CPU")
+  def test_example_readme_gpu(self): self._test_example_readme("GPU")
+  def test_example_readme_torch(self): self._test_example_readme("TORCH")
+  def test_example_readme_llvm(self): self._test_example_readme("LLVM")
+
+  def test_example_matmul_cpu(self): self._test_example_matmul("CPU")
+  def test_example_matmul_gpu(self): self._test_example_matmul("GPU")
+  def test_example_matmul_torch(self): self._test_example_matmul("TORCH")
+  def test_example_matmul_llvm(self): self._test_example_matmul("LLVM")
+
+if __name__ == '__main__':
+  unittest.main()