Skip to content

Subtensor lift: reason about advanced indices jointly when gating#2251

Open
ricardoV94 wants to merge 2 commits into
pymc-devs:mainfrom
ricardoV94:proper_idx_gating
Open

Subtensor lift: reason about advanced indices jointly when gating#2251
ricardoV94 wants to merge 2 commits into
pymc-devs:mainfrom
ricardoV94:proper_idx_gating

Conversation

@ricardoV94

@ricardoV94 ricardoV94 commented Jun 21, 2026

Copy link
Copy Markdown
Member

We had heuristics that work per-axis, but we should reason about the advanced indices together since they interact. E.g., tril_indices, would fail the per-axis checks, as each index (row, col) is larger than the respective axis and contains duplicates. But jointly they index less than the full array and don't have duplicates, so we want to lift it through Elemwise and the like.

Shows up in pymc-devs/pymc#8297

@ricardoV94 ricardoV94 marked this pull request as ready for review June 25, 2026 13:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant