Copy outputs in numba impl of Split

Author: ricardoV94

pytensor/link/numba/dispatch/tensor_basic.py

@@ -144,7 +144,8 @@ def split(tensor, axis, indices):
         # Work around for https://github.com/numba/numba/issues/8257
         axis = axis % tensor.ndim
         axis = numba_basic.to_scalar(axis)
-        return np.split(tensor, np.cumsum(indices)[:-1], axis=axis)
+        # TODO: Add inplace optimization for Split, so we don't need to copy all the time
+        return [a.copy() for a in np.split(tensor, np.cumsum(indices)[:-1], axis=axis)]
 
     return split
 

tests/link/numba/test_tensor_basic.py

@@ -4,10 +4,11 @@
 import pytensor.scalar as aes
 import pytensor.tensor as at
 import pytensor.tensor.basic as atb
-from pytensor import config
+from pytensor import config, function
 from pytensor.compile.sharedvalue import SharedVariable
 from pytensor.graph.basic import Constant
 from pytensor.graph.fg import FunctionGraph
+from pytensor.scalar import Add
 from pytensor.tensor.shape import Unbroadcast
 from tests.link.numba.test_basic import (
     compare_numba_and_py,
@@ -332,6 +333,30 @@ def test_Split(n_splits, axis, values, sizes):
     )
 
 
+def test_Split_not_inplace():
+    # https://github.com/pymc-devs/pytensor/issues/343
+    x1 = at.matrix("x1")
+    x2 = at.matrix("x2", shape=(None, 1))
+    v = at.vector("v", shape=(2,), dtype=int)
+    out = at.split(x1, v, n_splits=2, axis=1)[0] + x2
+
+    fn = function([x1, x2, v], out, mode="NUMBA")
+    # Check that the addition of split[0] and x2 is in place
+    add_op = fn.maker.fgraph.outputs[0].owner.op
+    assert isinstance(add_op.scalar_op, Add)
+    assert add_op.inplace_pattern == {0: 0}
+
+    rng = np.random.default_rng(123)
+    test_x1 = rng.normal(size=(2, 2))
+    test_x2 = rng.normal(size=(2, 1))
+    test_v = np.array([1, 1])
+
+    np.testing.assert_allclose(
+        fn(test_x1, test_x2, test_v).copy(),
+        fn(test_x1, test_x2, test_v).copy(),
+    )
+
+
 @pytest.mark.parametrize(
     "val, offset",
     [