pymc-devs / pymc-devs/pytensor
Create `vectorized`, `value_and_grad` and `shape` versions of `JAXOp`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
When we have a JAXOp in the final graph in a non-jax backend we may want to manipulate the JAX Op for efficiency. We could rewrite Blockwise(JAXOp) -> JAXOp whose inner function is vectorized.
If we have both the Op and the gradient, we could rewrite into a single op that uses value_and_grad under the hood.
And similarly if we only need the shape we could rewrite into an Op whose internal function only computes the shape. This last one is only relevant if the original Op doesn't remain in the graph.
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
Start by tracing the JAXOp and Blockwise rewrite paths described in the issue, then identify how the final graph represents an operation, its gradient, and shape computation. Define the three desired transformations—vectorized, value_and_grad, and shape—and determine how each can replace or combine the existing JAXOp without leaving the original operation when it is unnecessary.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100