Numba IndexedElemwise: Faster indexed increments when indices do not alias

[ricardoV94]

import numpy as np
import pytensor
import pytensor.tensor as pt
from pytensor.assumptions import assume

def make(use_assumption, n=8192):
    x = pt.vector("x")
    v = pt.vector("v")
    idx = pt.vector("idx", dtype="int64")
    idx_used = assume(idx, unique_indices=True) if use_assumption else idx
    val = pt.sqrt(v * v + 1.0)                       # vectorizable value-compute
    out = x[idx_used].inc(val)
    return pytensor.function([x, idx, v], out, trust_input=True)

n = 8192
rng = np.random.default_rng(0)
out0 = rng.random(n)
idxs = rng.permutation(n).astype("int64")            # unique -> the assumption holds
vals = rng.random(n)

fn_before = make(use_assumption=False)               # scatter RMW poisons the loop -> scalar
fn_after  = make(use_assumption=True)                # unique-index promise -> vectorizes

assert np.allclose(
   fn_before(out0.copy(), idxs, vals),
   fn_after(out0.copy(), idxs, vals),
)

%timeit fn_before(out0.copy(), idxs, vals)   # ~27 µs
%timeit fn_after(out0.copy(), idxs, vals)    # ~20 µs  (~1.4x)