pymc-devs / pymc-devs/pytensor
Rewrite nested inc/set_subtensor on zeros
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
The gradient of x[1:][-1] has two successive inc_subtensor on zeros of increasing size. We should collapse them, as happens if you take the gradient of the single slice that corresponds to the two slices together x[-1].
This shows up in the gradient of Scans for the last outputs of a recurring sequence.
import pytensor.tensor as pt
from pytensor.graph import rewrite_graph
x = pt.vector("x", shape=(4,))
out = x[1:][-1] # When you select the last entry of a scan sitsot this shows up in the graph
g = pt.grad(out, x)
rewrite_graph(g, include=("fast_run",), exclude=("inplace",)).dprint()
# IncSubtensor{start:} [id A]
# ├─ Alloc [id B]
# │ ├─ [0.] [id C]
# │ └─ 4 [id D]
# ├─ IncSubtensor{i} [id E]
# │ ├─ Alloc [id F]
# │ │ ├─ [0.] [id C]
# │ │ └─ 3 [id G]
# │ ├─ 1.0 [id H]
# │ └─ -1 [id I]
# └─ 1 [id J]
new_out = rewrite_graph(out, include=("ShapeOpt", "canonicalize"))
new_g = pt.grad(new_out, x)
rewrite_graph(new_g, include=("fast_run",), exclude=("inplace",)).dprint()
# IncSubtensor{i} [id A]
# ├─ Alloc [id B]
# │ ├─ [0.] [id C]
# │ └─ 4 [id D]
# ├─ 1.0 [id E]
# └─ 3 [id F]
I think the rule is incsubtensor on the larger buffer with the negative inner index, or outer start + positive inner index. We may also want to handle the unknown sign symbolically, but even the constant case would be a nice start.
Bonus points if we can combine it with an outer flip that the scan gradient also does:
import pytensor.tensor as pt
from pytensor.graph import rewrite_graph
x = pt.vector("x", shape=(4,))
out = x[1:][-1]
new_out = rewrite_graph(out, include=("ShapeOpt", "canonicalize"))
new_g = pt.grad(new_out, x)[::-1]
rewrite_graph(new_g, include=("fast_run",), exclude=("inplace",)).dprint()
# Subtensor{::step} [id A]
# ├─ IncSubtensor{i} [id B]
# │ ├─ Alloc [id C]
# │ │ ├─ [0.] [id D]
# │ │ └─ 4 [id E]
# │ ├─ 1.0 [id F]
# │ └─ 3 [id G]
# └─ -1 [id H]
Which should be doable by flipping the indices. Not as important since the flip is just a cheap view on the input
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 running the provided pt.grad and rewrite_graph examples and tracing the rewrite rules that handle IncSubtensor and nested slice indices. Done means the two IncSubtensor operations on zero buffers collapse for the constant-index cases shown; the outer flip combination is optional bonus work.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100