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 147 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 ⚠️

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%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[jessegrabowski]

Some typos and questions