Do not use Numba objmode for supported advanced indexing operations

[ricardoV94]

Description

Several cases of advanced indexing are supported by Numba these days. This PR makes sure we don't skip on those!

Related Issue

• Closes #
• Related to Handle more AdvancedSetSubtensor ops in numba #772

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

[codecov]

Codecov Report

Attention: Patch coverage is 94.78261% with 6 lines in your changes missing coverage. Please review.

Project coverage is 80.87%. Comparing base (36c55f5) to head (b6ee8e7).
Report is 191 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/numba/dispatch/subtensor.py	94.39%	3 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #805      +/-   ##
==========================================
+ Coverage   80.76%   80.87%   +0.11%     
==========================================
  Files         163      164       +1     
  Lines       46941    46891      -50     
  Branches    11499    11474      -25     
==========================================
+ Hits        37912    37925      +13     
+ Misses       6786     6752      -34     
+ Partials     2243     2214      -29     

Files with missing lines	Coverage Δ
pytensor/link/numba/dispatch/__init__.py	100.00% <100.00%> (ø)
pytensor/link/numba/dispatch/basic.py	86.35% <ø> (+0.80%)	⬆️
pytensor/tensor/subtensor.py	89.89% <100.00%> (+0.02%)	⬆️
pytensor/tensor/type_other.py	74.15% <100.00%> (-8.99%)	⬇️
pytensor/link/numba/dispatch/subtensor.py	94.39% <94.39%> (ø)

... and 12 files with indirect coverage changes

[ricardoV94]

@kc611 sorry for tagging if your review, fat fingers on my side! Do keep the cool work on extending advanced indexing on the Numba side!!!

[kc611]

@kc611 sorry for tagging if your review, fat fingers on my side! Do keep the cool work on extending advanced indexing on the Numba side!!!

No worries! 🙂 I took a quick look at this PR and I'd say if you guys are using string-based function creation anyways then it's completely possible to support the full set of advanced indexing cases by simply creating the required loops within the string-based functions.

[ricardoV94]

@kc611 sorry for tagging if your review, fat fingers on my side! Do keep the cool work on extending advanced indexing on the Numba side!!!

No worries! 🙂 I took a quick look at this PR and I'd say if you guys are using string-based function creation anyways then it's completely possible to support the full set of advanced indexing cases by simply creating the required loops within the string-based functions.

Yeah, but we are not doing that yet. We are just creating dumb stuff like x[indices] = y. We'll come to that if we need it, or perhaps Numba will support it first :D

Thanks for checking it out!

[ricardoV94]

This should be ready for review. After this we can start adding some simple non-supported cases like boolean indexing with indices.shape == x.shape, and some multiple vector index cases

[jessegrabowski]

Looks good to me. The duplicate indices thing isn't really clear to me. Also curious if the lines codecov is complaining about are unreachable, and, if so, whether we can just get rid of them.

[jessegrabowski]

Here "default" means on our end? Is that a problem?

[ricardoV94]

It means the code returned by numba_funcify_default_subtensor. Shouldn't be a problem as long as we don't try to use it when it's not valid