Support more cases of numba advanced indexing

[ricardoV94]

Support more cases of multi-dimensional advanced indexing and updating in Numba

Extends pre-existing rewrite to ravel multidimensional integer indices, to handle multiple inputs (if no broadcasting is needed) and to place them consecutively if they were not. It also extends it to Set/IncSubtensor as long as y is not broadcasted and advanced indices are consecutive.

The following cases should now be supported without object mode:

• Advanced integer indexing (not mixed with basic or boolean indexing) that do not require broadcasting of indices
• Consecutive advanced integer indexing updating (set/inc) (not mixed with basic or boolean indexing) that do not require broadcasting of indices or y.

Also fixes bug in infer_shape of AdvancedIndexing with slices (which were mistakenly treated as NoneSlices)

Example of new kind of graphs that are supported.

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

x = pt.tensor("x", shape=(10, 10, 3))
inds = pt.matrix("inds", shape=(50, 2), dtype=int)
y = x[inds, inds]
x_test = np.zeros((10, 10, 3)) 
inds_test = np.ones((50, 2), dtype=int)

fn = pytensor.function([x, inds], [y, pt.grad(y.sum(), x)], mode="NUMBA")
fn.dprint(print_shape=True)
fn.trust_input = True
print(fn(x_test, inds_test)[0].shape)
%timeit fn(x_test, inds_test)

On my machine that runs in 8us, and before with object mode it was 60us.

A case with a single matrix advanced indexing shows up in the logp of the Categorical, the gradient of which was not supported by numba without object mode before and now is.

---

📚 Documentation preview 📚: https://pytensor--1254.org.readthedocs.build/en/1254/

[codecov]

Codecov Report

Attention: Patch coverage is 86.95652% with 9 lines in your changes missing coverage. Please review.

Project coverage is 81.99%. Comparing base (2a7f3e1) to head (d1f18a1).
Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/subtensor.py	77.77%	5 Missing and 1 partial ⚠️
pytensor/tensor/rewriting/subtensor.py	92.85%	1 Missing and 2 partials ⚠️

❌ Your patch status has failed because the patch coverage (86.95%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1254      +/-   ##
==========================================
- Coverage   81.99%   81.99%   -0.01%     
==========================================
  Files         188      188              
  Lines       48553    48600      +47     
  Branches     8673     8685      +12     
==========================================
+ Hits        39812    39849      +37     
- Misses       6579     6586       +7     
- Partials     2162     2165       +3     

Files with missing lines	Coverage Δ
pytensor/link/numba/dispatch/subtensor.py	95.34% <ø> (ø)
pytensor/tensor/rewriting/subtensor.py	90.12% <92.85%> (-0.13%)	⬇️
pytensor/tensor/subtensor.py	89.35% <77.77%> (-0.37%)	⬇️

[jessegrabowski]

Some typos and questions