Add general C kernels for AdvancedSubtensor and AdvancedIncSubtensor

[velochy]

Description

AdvancedSubtensor and AdvancedIncSubtensor only had C implementations for two narrow patterns: a single integer index with every other axis in full (via PyArray_TakeFrom), and x[idx] (+)= y with 1-D x. Everything else ran through perform, i.e. x.__getitem__ and np.add.at, which is where the time goes in a C-linked function as soon as a model indexes with two arrays, updates a 2-D parameter by row, or takes a gradient through x[idx].

This replaces both with one code generator, _AdvIndexCGen, that emits a statically unrolled loop nest for the general integer case:

• any number of integer index arrays, of any ndim, broadcast against each other (0-d indices included);
• mixed with full slices in any position, consecutive or not (the index group moves to the front in the non-consecutive case, as in numpy);
• any stride layout for x, y and the indices;
• bounds checks and negative-index wrapping inline, with numpy's error messages.

The loops run in the memory order of the indexed shape, so the output of a gather (and y of a scatter) is walked sequentially. The innermost loop takes a memcpy when it is a same-type copy over unit strides and an indexable, vectorizable loop otherwise. Dims and strides are hoisted into const locals up front so stores through char* cannot force reloads inside the loops. There is no per-call allocation.

Still on perform: boolean masks, non-full slices inside the advanced op, ignore_duplicates=True (numba keeps numpy's last-write-wins there, which a sequential scatter can't), complex dtypes, and += into a bool array.

y of the scatter may broadcast via static 1-dims, missing leading dims, or as a scalar. A runtime length-1 dim that is not static-broadcastable raises the same "Runtime broadcasting not allowed" error the old 1-D C path raised; perform only checks this for the vector-index case.

Timings

ms per call, float64, int64 indices, 20000 index positions, x is (200, 50) or (20, 30, 40). Min of 5 interleaved rounds on a Ryzen 4750U.

case	CVM before	CVM after	Numba
x[i, j]	0.138	0.053	0.047
x[i, j] += y	0.563	0.069	0.054
x[i]	2.22	2.20	1.84
x[i] += Y	23.6	1.36	1.24
x[i].set(Y)	2.48	1.60	1.39
x[:, j]	8.09	7.54	8.35
x[:, j] += Y	111.6	8.22	9.61
x3[i, :, k]	2.05	1.25	1.28
x3[i, :, k] += Y	15.5	1.13	1.15
x3[:, i]	40.0	40.6	118
x3[:, i] += Y	424	15.0	36.8
x[I2, J2] (2-D indices)	0.143	0.057	0.047
x[I2, J2] += Y	0.578	0.091	0.056

The two gathers that already had a C path (x[i], x3[:, i]) are unchanged within noise; both are bound by writing the 8 MB / 128 MB output.

Tests

TestAdvancedIndexingCImpl compares the cvm and py linkers, asserting the C path is taken, over 15 index patterns × C/F-ordered x for gathers and × inc/set for scatters (including y broadcasting), plus dtype mixing (float32 += float64, int32 += float64, bool set from float) and the error paths (out of bounds, index shape mismatch, y shape mismatch, runtime broadcast). tests/tensor/test_subtensor.py and the subtensor rewrite tests pass under linker=cvm.

Related Issue

• Closes #
• Related to Remove AdvancedSubtensor1 and AdvancedIncSubtensor1 Ops #2294

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):

🤖 Generated with Claude Code

[velochy]

@ricardoV94 I know CVM is no longer the default so optimizing it is not high priority, but I do have a use case. Namely, numba lacks good profiling tools so whenever I ask AI to try to figure out computational bottlenecks, it uses CVM as proxy - and, invariably first flags indexing as a major source, until it benchmarks with numba to confirm it is not.

This PR brings indexing performance to parity with numba in most cases, hopefully making future profiling proxy runs skip that step.

[ricardoV94]

CVM is very different than numba not just indexing (per op dispatch, no major fusion besides Elemwise).

When I want to evaluate numba code quality I usually:

1. tell bot to look at emitted llvm
2. tell bot to write the fastest direct numba implementation it can (sidestep pytensor altogether).

Then make use of the two sources to look for obvious inefficiencies.

[ricardoV94]

This PR brings indexing performance to parity with numba in most cases, hopefully making future profiling proxy runs skip that step.

I doubt that because these days we rarely do indexing for indexing sake in the numba backend. we either fuse with the elemwise consumer (gather) or producer (scatter)