Skip to content

Conversation

@ricardoV94
Copy link
Member

(and mixed basic indexing)

This covers the cases with default ignore_updates=False for inc_subtensor,

Boolean indices are converted to integer indices during specialize rewriting.

This was motivated by #811 where it found the logic for when object fallback was needed was wrong.
Fixing it, revealed we had less cases covered by pure numba than thought, hence this PR.

Copy link

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR completes support for integer advanced indexing in PyTensor's Numba backend by implementing a comprehensive dispatch mechanism that handles all forms of advanced indexing with basic and vector integer indices, including non-consecutive indices and mixed basic/advanced indexing patterns.

Key Changes:

  • Replaces specialized multidimensional bool/int index raveling rewrites with a simpler bool_idx_to_nonzero rewrite that converts boolean indices to integer indices via nonzero()
  • Implements a new codegen-based vector_integer_advanced_indexing dispatcher that generates custom Numba functions for each indexing pattern, eliminating most object mode fallbacks
  • Updates tests to remove the objmode_needed parameter and consolidates test parameters for ignore_duplicates behavior

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 7 comments.

File Description
tests/link/numba/test_subtensor.py Removes objmode_needed parameter from tests and consolidates duplicate index test parameters, reflecting that most indexing operations no longer require object mode fallback
pytensor/tensor/rewriting/subtensor.py Simplifies boolean indexing rewrite by converting to nonzero() instead of raveling, and removes complex multidimensional integer indexing rewrites that are now handled by the dispatcher
pytensor/link/numba/dispatch/subtensor.py Implements comprehensive codegen-based dispatcher for advanced indexing that handles transpose operations, broadcasting, and index raveling at runtime without requiring object mode

@ricardoV94 ricardoV94 force-pushed the numba_advanced_vec_indexing branch 2 times, most recently from 59b83a4 to ebef3b1 Compare December 9, 2025 13:50
…) advanced indexing

When default `ignore_updates=True` for inc_subtensor, and boolean indices were rewritten during specialize
@ricardoV94 ricardoV94 force-pushed the numba_advanced_vec_indexing branch from ebef3b1 to 61a9309 Compare December 9, 2025 14:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

backend compatibility bug Something isn't working enhancement New feature or request indexing numba

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant