pymc-devs / pymc-devs/pytensor
AdvancedSubtensor1 and AdvancedIncSubtensor1 are not needed in JAX and Numba backends
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
AdvancedSubtensor1 and AdvancedIncSubtensor1 are specializations of AdvancedSubtensor and AdvancedIncSubtensor that only allow one vector of integer indexing along the first dimension. They exist to provide a C-implementation of this more restricted form of advanced indexing (the other Ops just use Python code), with some rewrites that try to convert from general form by transposing/reshaping stuff back and forth.
These shouldn't be needed in the JAX/Numba backends, because those backends already support a single dimension index wherever it may be, so no need for the rewrite that introduces those in the backends?
Well AdvancedIncSubtensor1 may still be fine in the Numba backend for now, we have some specialized implementations for those.
We should check we have no other useful rewrites that depend on these specialized Ops
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Trace where AdvancedSubtensor1 and AdvancedIncSubtensor1 are introduced by rewrites for the JAX and Numba backends. Check whether other useful rewrites depend on these specialized Ops, while preserving the Numba-specific AdvancedIncSubtensor1 implementations mentioned in the issue. Done means the unnecessary backend rewrites are removed or retained with clear justification and relevant backend behavior still works.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, compilers
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 30/100