Numba FillDiagonal: Do not mutate input

Author: ricardoV94

pytensor/link/numba/dispatch/extra_ops.py

@@ -100,16 +100,19 @@ def cumop(x):
 def numba_funcify_FillDiagonal(op, **kwargs):
     @numba_basic.numba_njit
     def filldiagonal(a, val):
+        a = a.copy()
         np.fill_diagonal(a, val)
         return a
 
-    return filldiagonal
+    cache_version = 1
+    return filldiagonal, cache_version
 
 
 @register_funcify_default_op_cache_key(FillDiagonalOffset)
 def numba_funcify_FillDiagonalOffset(op, node, **kwargs):
     @numba_basic.numba_njit
     def filldiagonaloffset(a, val, offset):
+        a = a.copy()
         height, width = a.shape
         offset_item = offset.item()
         if offset >= 0:
@@ -128,7 +131,8 @@ def filldiagonaloffset(a, val, offset):
         # return a
         return b.reshape(a.shape)
 
-    return filldiagonaloffset
+    cache_version = 1
+    return filldiagonaloffset, cache_version
 
 
 @register_funcify_default_op_cache_key(RavelMultiIndex)

tests/link/numba/test_extra_ops.py

@@ -84,7 +84,6 @@ def test_CumOp(val, axis, mode):
     )
 
 
-@pytest.mark.xfail(reason="Implementation works inplace!")
 def test_FillDiagonal():
     a = pt.lmatrix("a")
     test_a = np.zeros((10, 2), dtype="int64")