Skip to content

Correct nit_sot shapes for scan with n_steps=0 (or raise)#2040

Draft
ricardoV94 wants to merge 8 commits intopymc-devs:v3from
ricardoV94:scan_nitsot_hack
Draft

Correct nit_sot shapes for scan with n_steps=0 (or raise)#2040
ricardoV94 wants to merge 8 commits intopymc-devs:v3from
ricardoV94:scan_nitsot_hack

Conversation

@ricardoV94
Copy link
Copy Markdown
Member

@ricardoV94 ricardoV94 commented Apr 10, 2026

When n_steps=0, Scan's perform method was creating nit_sot output arrays with shape (0,)*ndim instead of (0, *core_shape). This caused silent shape mismatches in downstream operations like concatenation or broadcasting.

Add Scan.nit_sot_core_shapes() method that returns static core shapes when all are known, or None otherwise. All three backends (Python, Cython, Numba) now use correct shapes when known, and raise a clear
ValueError when shapes are unknown and n_steps=0.

We could consider making the shapes a symbolic input (PyTensor can often lift it out of the computation), a bit like numba BlockwiseWithCoreShape and RandomVariableWithCoreShape, but that's an improvement, not a must.

Now we don't silently return garbage.

Closes #1287

Rename Op.grad/L_op/R_op to Op.pullback/pushforward with deprecation
wrappers. Rename module-level Rop/Lop to pushforward/pullback with
deprecation wrappers. Rename _rop_legacy to _pushforward_direct.
Update parameter names: eval_points to tangents/cotangents,
use_op_rop_implementation to use_op_pushforward.
@ricardoV94 ricardoV94 changed the title Fix Scan n_steps=0 producing wrong nit_sot output shapes Correct nit_sot shapes for scan with n_steps=0 (or raise) Apr 10, 2026
jessegrabowski and others added 3 commits April 10, 2026 14:53
When n_steps=0, Scan's perform method was creating nit_sot output
arrays with shape (0,)*ndim instead of (0, *core_shape). This caused
silent shape mismatches in downstream operations like concatenation
or broadcasting.

Add Scan.nit_sot_core_shapes() method that returns static core shapes
when all are known, or None otherwise. All three backends (Python,
Cython, Numba) now use correct shapes when known, and raise a clear
ValueError when shapes are unknown and n_steps=0.
@review-notebook-app
Copy link
Copy Markdown

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Scan NIT-SOT with 0 steps have wrong shape

2 participants