Fix non-inplace IfElse on numba backend

[emekaokoli19]

Description

This PR fixes an issue in the IfElse numba, where outputs were returned as direct references to the input arrays instead of copies. This violated the semantics of ifelse, which guarantees that the returned value is a distinct object, even when both branches reference the same input.

The fix ensures that each selected output is explicitly copied. This matches the behavior of the Python linker and prevents unexpected mutations when the NumPy arrays are assumed to be non-shared.

Related Issue

• Closes Numba ifelse does not copy inputs when not inplace #1761
• Related to #

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[ricardoV94]

Nice trick with the list, I was thinking we would need to use codegen due to numba limitations.

Left some comments.

We also need to work on the existing tests so they would have failed before the fix and pass now. We have to test both single and multi output and inplace or not

[ricardoV94]

The list trick didn't work. I suspect that would happen.

We need to use codegen for the non inplace version of ifelse with multiple outputs. Other cases will work fine.

You can see some cases of codegen for the numba funcify Alloc or SpecifyShape that may give some ideas. This tends to happen everytime we would want (*args) in the signature of a numba func

[emekaokoli19]

The list trick didn't work. I suspect that would happen.

We need to use codegen for the non inplace version of ifelse with multiple outputs. Other cases will work fine.

You can see some cases of codegen for the numba funcify Alloc or SpecifyShape that may give some ideas. This tends to happen everytime we would want (*args) in the signature of a numba func

Hey @ricardoV94, I have added codegen. A lot of tests are failing in tests/link/numba, so I wrote some new tests to test the ifelse changes. Are the failing tests in tests/link/numba a result of something else?

[emekaokoli19]

@ricardoV94

Explanation for failing is checks in Numba mode tests

The tests use assert res is a to check if the output is the same Python object as the input. This fails under Numba and I think it is because:

1. IfElse with as_view=True normally returns the same object in Python mode.
2. In Numba mode, the compiled function creates new arrays for outputs, even when as_view=True.
3. Therefore, res is a fails because object identity is different, even though the array values are identical.

Should we replace is checks with np.array_equal, which tests element-wise equality instead of memory identity? However, this would mean that as_view would be redundant with numba mode

[ricardoV94]

Ah you need to set the borrow flags to allow pytensor to pass back inputs unalterated:

import numpy as np
import pytensor
import pytensor.tensor as pt

x = pt.vector("x")
fn = pytensor.function([pytensor.In(x, borrow=True)], pytensor.Out(x, borrow=True), mode="NUMBA")
x_test = np.zeros(5)
fn(x_test) is x_test

You can check that without that the final function would have a deepcopy Op, using fn.dprint()

[emekaokoli19]

@ricardoV94, Firstly, I am sorry for the delay on this PR. I have been out sick

I think the test is failing because Python object identity is lost at the graph level, before execution reaches Numba.
This is not caused by Numba. It is caused by a Subtensor{i} → ViewOp chain applied to the output of IfElse.

Observed graph from the test is below

From fn.dprint():

ViewOp [id A]
 └─ Subtensor{i} [id B]
    ├─ if{inplace} [id C]
    │  ├─ Gt
    │  │  ├─ Sum{axes=None}
    │  │  │  └─ x
    │  │  └─ 0
    │  ├─ ExpandDims{axis=0}
    │  │  └─ x
    │  └─ ExpandDims{axis=0}
    │     └─ x
    └─ 0

Subtensor always creates a new NumPy array object.

[ricardoV94]

Subtensor should always return a view. You can add print_view_map=True to dprint to see the information on views.

But why do you have a subtensor? Your [0] is not selecting the first output, is selecting the first entry of the first output. I think when IfElse has a single variable in each branch, it doesn't return a list with it, just returns the variable unless you pass return_list=True

[ricardoV94]

I (force-)pushed some changes. This bug was causing pymc-devs/pymc#7993 to segfault.

I also found that IfElse never views false branch variables, only true, due to a limitation in what PyTensor can handle for inplace

got involved in code changes

[jessegrabowski]

This comment is too cryptic; can you put a small note why they aren't aliased? (And by "aliased" do you mean "views"?) You mentioned "a limitation in pytensor" in a comment on this PR but I really don't get it.

[ricardoV94]

View_map is always output: input_from_true_branch. More context in #1811

[ricardoV94]

Changed to # We only view true branch variables. False branch variables must always be copied.

# We only ever view (alias) variables from the true branch. False branch variables must always be copied.

More clear?

[jessegrabowski]

It's more clear but I still don't understand why. It seems strange and arbitrary.

[ricardoV94]

It's how the Op is implemented. The limitation is part of the current inplace limitations that go beyond IfElse, as mentioned in #1811