Add numba impl for CheckAndRaise

Author: aseyboldt

pytensor/link/numba/dispatch/extra_ops.py

@@ -19,6 +19,7 @@
     Unique,
     UnravelIndex,
 )
+from aesara.raise_op import CheckAndRaise
 
 
 @numba_funcify.register(Bartlett)
@@ -372,3 +373,18 @@ def broadcast_to(x, *shape):
         return np.broadcast_to(x, scalars_shape)
 
     return broadcast_to
+
+
+@numba_funcify.register(CheckAndRaise)
+def numba_funcify_CheckAndRaise(op, node, **kwargs):
+    error = op.exc_type
+    msg = op.msg
+
+    @numba_basic.numba_njit
+    def check_and_raise(x, *conditions):
+        for cond in conditions:
+            if not cond:
+                raise error(msg)
+        return x
+
+    return check_and_raise