JitError (#14041)

* JitError

* test_symbolic_jit

test/test_jit.py

@@ -5,7 +5,7 @@ import numpy as np
 from hypothesis import given, settings, strategies as strat
 from test.helpers import assert_jit_cache_len, not_support_multi_device, REAL_DEV, needs_second_gpu
 from tinygrad.tensor import Tensor
-from tinygrad.engine.jit import TinyJit, GraphRunner, MultiGraphRunner, graph_class
+from tinygrad.engine.jit import TinyJit, JitError, GraphRunner, MultiGraphRunner, graph_class
 from tinygrad.engine.realize import CompiledRunner, BufferCopy, BufferXfer
 from tinygrad.device import Device
 from tinygrad.helpers import Context, JIT, GlobalCounters, getenv
@@ -76,7 +76,7 @@ class TestJit(unittest.TestCase):
   def test_nothing_jitted(self):
     @TinyJit
     def add(a, b): return None
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(JitError):
       for _ in range(5):
         a = Tensor.randn(10, 10)
         b = Tensor.randn(10, 10)
@@ -125,13 +125,13 @@ class TestJit(unittest.TestCase):
       b = Tensor.randn(10, 10)
       add(a, b)
     bad = Tensor.randn(20, 20)
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(JitError):
       add(a, bad)
 
   def test_jit_shape_views_mismatch(self):
     @TinyJit
     def add(a): return (a+1).realize()
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(JitError):
       for i in range(1,5):
         # a has an offset that the kernel doesn't know about
         a = Tensor.randn(10, 10).realize()[:, i:i+2]
@@ -142,7 +142,7 @@ class TestJit(unittest.TestCase):
     @TinyJit
     def add(a, b): return (a+b).realize()
     a = Tensor.randn(10, 10)
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(JitError):
       add(a, a)
 
   def test_jit_assign(self, dtype=dtypes.float32):
@@ -510,7 +510,7 @@ class TestJit(unittest.TestCase):
     # TODO: this should fail since input has a different size
     f(Tensor(2.0)).item()
     # TODO: this should not fail, and should return 3
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(JitError):
       f(Tensor([2.0])).item()
 
 @unittest.skip("Pending multioutput implementation #3607")

test/test_jit_footguns.py

@@ -21,6 +21,7 @@ ERRORS RAISED (lower priority - at least users know):
 import unittest
 import numpy as np
 from tinygrad import Tensor, TinyJit
+from tinygrad.engine.jit import JitError
 
 class TestJitFootguns(unittest.TestCase):
 
@@ -70,7 +71,7 @@ class TestJitFootguns(unittest.TestCase):
     def f(a, b): return (a + b).realize()
 
     x = Tensor([1, 2, 3])
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(JitError):
       f(x, x)
 
   def test_tensors_in_containers_ignored(self):
@@ -116,7 +117,7 @@ class TestJitFootguns(unittest.TestCase):
     def f(a): return (a + 1).realize()
 
     base = Tensor.randn(10, 10).realize()
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(JitError):
       for i in range(1, 5):
         f(base[:, i:i+2])  # different offset each time
 
@@ -128,7 +129,7 @@ class TestJitFootguns(unittest.TestCase):
     f(Tensor.randn(10, 10), Tensor.randn(10, 10))  # warmup
     f(Tensor.randn(10, 10), Tensor.randn(10, 10))  # capture
 
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(JitError):
       f(Tensor.randn(20, 20), Tensor.randn(20, 20))
 
   def test_python_constants_frozen(self):
@@ -170,7 +171,7 @@ class TestJitFootguns(unittest.TestCase):
     f(Tensor([1]), Tensor([2]))  # warmup with positional
     f(Tensor([1]), Tensor([2]))  # capture with positional
 
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(JitError):
       f(a=Tensor([3]), b=Tensor([4]))  # kwargs fail
 
   def test_class_method_shared_across_instances(self):
@@ -213,7 +214,7 @@ class TestJitFootguns(unittest.TestCase):
     @TinyJit
     def f(a, b): return None
 
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(JitError):
       for _ in range(3):
         f(Tensor([1]), Tensor([2]))
 

test/test_symbolic_jit.py

@@ -2,6 +2,7 @@ import unittest
 
 from test.helpers import assert_jit_cache_len
 from tinygrad import Variable, Tensor, TinyJit
+from tinygrad.engine.jit import JitError
 import numpy as np
 
 class TestSymbolicJit(unittest.TestCase):
@@ -172,7 +173,7 @@ class TestSymbolicJit(unittest.TestCase):
     vi2 = Variable("i", 1, 10).bind(7)
     a = Tensor.rand(3, 7)[:, :vi2]
     bad = Tensor.rand(4, 7)[:, :vi2]
-    with self.assertRaises(AssertionError):
+    with self.assertRaises(JitError):
       add(a, bad)
 
   def test_shrink(self):