Numba: speed up advanced indexing with non-scalar core dims

[ricardoV94]

vector_integer_advanced_indexing emitted one whole-subarray operation per index step (out_buffer[i] = basic_indexed_x[scalar_idxs], and the += form for scatters). Numba routes a subarray assignment through generic array machinery, and the += form materialises a temporary per step.

Emit an explicit scalar loop nest over the trailing dimensions instead.

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

rng = np.random.default_rng(0)
x_val = rng.normal(size=(1000, 5))
idx_val = rng.integers(0, 1000, 10_000)
y_val = rng.normal(size=(10_000, 5))

x, y = pt.matrix("x"), pt.matrix("y")
idx = pt.lvector("idx")

take = pytensor.function([x, idx], x[idx], mode="NUMBA", trust_input=True)
add_at = pytensor.function([x, y, idx], x[idx].inc(y), mode="NUMBA", trust_input=True)
take(x_val, idx_val), add_at(x_val, y_val, idx_val)  # warm up compilation

%timeit take(x_val, idx_val)
%timeit add_at(x_val, y_val, idx_val)

Before:

178 μs ± 1.52 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
319 μs ± 28.6 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

After:

43.1 μs ± 1.66 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
69.6 μs ± 6.68 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

I was shocked that we were leaving so much on the floor tbh

[ricardoV94]

CC @velochy (one of four implicated PRs on the MRP model catch-up)

[velochy]

My guess is this will be huge for us. We have a lot of indexing in our main flows, not just for likerts but for our order models as well. Very possible this will end up halving our execution time on the latter, which would be a pretty big deal

[velochy]

pl_bestworst_repro.py
For reference: this is a pared-down version of our order models. Currently 1.3x vs jax, but my guess is this PR will take it to considerably below.

LMK if you would like this to be in the pymc-examples as well

[ricardoV94]

let me first collect the changes and do a release, then we can check with other variants