Merge branch 'main' into fix_dimshuffle

aesara-devs · Oct 22, 2021 · cda0357 · cda0357
2 parents d4dd922 + 0f8c81c
commit cda0357
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 0 deletions.
diff --git a/aesara/link/numba/dispatch/scalar.py b/aesara/link/numba/dispatch/scalar.py
@@ -21,6 +21,7 @@
     Clip,
     Composite,
     Identity,
+    Inv,
     Mul,
     ScalarOp,
     Second,
@@ -170,3 +171,12 @@ def second(x, y):
         return y
 
     return second
+
+
+@numba_funcify.register(Inv)
+def numba_funcify_Inv(op, node, **kwargs):
+    @numba.njit(inline="always")
+    def inv(x):
+        return 1 / x
+
+    return inv
diff --git a/tests/link/test_numba.py b/tests/link/test_numba.py
@@ -804,6 +804,25 @@ def test_Cast(v, dtype):
     )
 
 
+@pytest.mark.parametrize(
+    "v, dtype",
+    [
+        (set_test_value(aet.iscalar(), np.array(10, dtype="int32")), aesb.float64),
+    ],
+)
+def test_Inv(v, dtype):
+    g = aesb.inv(v)
+    g_fg = FunctionGraph(outputs=[g])
+    compare_numba_and_py(
+        g_fg,
+        [
+            i.tag.test_value
+            for i in g_fg.inputs
+            if not isinstance(i, (SharedVariable, Constant))
+        ],
+    )
+
+
 @pytest.mark.parametrize(
     "v, shape, ndim",
     [