pymc-devs / pymc-devs/pytensor
Scan save mem rewrite masks issues with `steps=0`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
import pytensor
import pytensor.tensor as pt
from pytensor.compile.mode import get_default_mode
n = pt.iscalar("n")
x0 = pt.vector("x0")
xs, _ = pytensor.scan(lambda xtm1: xtm1 + 1, outputs_info=[x0], n_steps=n)
out = xs[-1] # Invalid when nsteps=0
fn = pytensor.function([n, x0], out)
print(fn(n=0, x0=[0, 1])) # [1. 2.]
fn = pytensor.function([n, x0], out, mode=get_default_mode().excluding("shape_unsafe"))
print(fn(n=0, x0=[0, 1])) # [1. 2.]
fn = pytensor.function([n, x0], out, mode=get_default_mode().excluding("scan_save_mem"))
print(fn(n=0, x0=[0, 1])) # IndexError: index out of bounds
I suspect from this hack: https://github.com/pymc-devs/pytensor/blob/8454c3ba78b1889987a909745f87c8dcab8e48fb/pytensor/scan/rewriting.py#L1438-L1445
But removing this hack leads to some tests failing, so other stuff may be doing wrong assumptions downstream of it (or perhaps inside the rewrite)
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 with the reproducer in the issue and inspect pytensor/scan/rewriting.py around lines 1438-1445. Compare behavior with the default mode, shape_unsafe excluded, and scan_save_mem excluded, then run the tests that fail when the hack is removed. Done means steps=0 no longer produces an incorrect result while the existing scan rewrite tests continue to pass.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100