diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5bf60083f9..ea2a2053ed 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/test/unit/test_utils.py b/test/extra/test_utils.py similarity index 100% rename from test/unit/test_utils.py rename to test/extra/test_utils.py diff --git a/test/test_example.py b/test/test_example.py deleted file mode 100644 index 26c9a538b1..0000000000 --- a/test/test_example.py +++ /dev/null @@ -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() diff --git a/test/unit/test_example.py b/test/unit/test_example.py new file mode 100644 index 0000000000..63bffa502d --- /dev/null +++ b/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()