Implement Numba conversion for ExtractDiag

aesara/link/numba/dispatch.py

@@ -35,6 +35,7 @@ from aesara.tensor.basic import (
     AllocDiag,
     AllocEmpty,
     ARange,
+    ExtractDiag,
     Join,
     MakeVector,
     Rebroadcast,
@@ -819,3 +820,16 @@ def numba_funcify_Join(op, **kwargs):
         return np.concatenate(tensors, to_scalar(axis))
 
     return join
+
+
+@numba_funcify.register(ExtractDiag)
+def numba_funcify_ExtractDiag(op, **kwargs):
+    offset = op.offset
+    # axis1 = op.axis1
+    # axis2 = op.axis2
+
+    @numba.njit
+    def extract_diag(x):
+        return np.diag(x, k=offset)
+
+    return extract_diag

tests/link/test_numba.py

@@ -1021,3 +1021,32 @@ def test_Join_view():
                 if not isinstance(i, (SharedVariable, Constant))
             ],
         )
+
+
+@pytest.mark.parametrize(
+    "val, offset",
+    [
+        (
+            set_test_value(
+                aet.matrix(), np.arange(10 * 10, dtype=config.floatX).reshape((10, 10))
+            ),
+            0,
+        ),
+        (
+            set_test_value(aet.vector(), np.arange(10, dtype=config.floatX)),
+            0,
+        ),
+    ],
+)
+def test_ExtractDiag(val, offset):
+    g = aet.diag(val, offset)
+    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))
+        ],
+    )