Avoid default allocation for taps of length 1 in ScanSaveMem

[ricardoV94]

The check we had for whether a variable was a default scan buffer (that's not also broadcasting the initval), always failed for single tapped outputs. There's a conservative check that the original value is not being broadcast to the number of initial taps, but that doesn't matter for single taps.

Also added some checks that we are actually only keeping buffers of the expected size in the test.

---

📚 Documentation preview 📚: https://pytensor--1395.org.readthedocs.build/en/1395/

[Copilot]

Pull Request Overview

This PR addresses an issue with the default scan buffer allocation for single-tapped outputs in ScanSaveMem and enhances the tests for buffer size validation. Key changes include:

• Adjusting the test configuration by excluding "scan_pushout" and renaming an internal function from f_rnn to step for clarity.
• Updating the implementation of default scan buffer handling by adding a new parameter (taps) to _is_default_scan_buffer and adapting buffer expansion and slicing logic accordingly.

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
tests/scan/test_rewriting.py	Modified test configuration and assertions regarding scan buffer sizes and function naming.
pytensor/scan/rewriting.py	Updated _is_default_scan_buffer's signature and revised buffer handling logic using the taps value.

[ricardoV94]

this would fail with the infamous [-0:] edge case

[ricardoV94]

Only a minor float32 tolerance failure, so it's ready for review while I don't come and tweak it.

[jessegrabowski]

I don't get it, why are we expanding dims then immediately discarding the extra dims

[ricardoV94]

It may add more implicit dims, and we only want to discard the first one

[jessegrabowski]

but [0] only selects the first one and discards all the rest?

[ricardoV94]

No, this is a tensorvariable, it selects into something with ndim-1 dimensions, it's basically a squeeze(x, 0) which I can use it it's more readable

[lucianopaz]

If I'm reading this right, it's keeping all extra dims after discarding the first one. But I don't know enough of the inner workings of the Scan Op to understand what the two inputs are and how they interact

[ricardoV94]

Added a comment and switched to squeeze on the initval

[jessegrabowski]

Is this the default buffer case? (just uninitialized)

[ricardoV94]

default buffer is an empty with enough space to hold all scan steps plus initial taps, initial taps are written to the beginning of the buffer with set_subtensor

[jessegrabowski]

so is this the default case?

[ricardoV94]

I don't understand the question, if it's not an AllocEmpty it's not a default buffer, hence why it returns False. Am I missing something?

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 82.12%. Comparing base (2d414d4) to head (cb61b0c).
Report is 5 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1395   +/-   ##
=======================================
  Coverage   82.12%   82.12%           
=======================================
  Files         211      211           
  Lines       49722    49777   +55     
  Branches     8820     8830   +10     
=======================================
+ Hits        40832    40881   +49     
- Misses       6710     6715    +5     
- Partials     2180     2181    +1     

Files with missing lines	Coverage Δ
pytensor/scan/rewriting.py	82.92% <100.00%> (+0.02%)	⬆️
pytensor/tensor/rewriting/basic.py	94.59% <100.00%> (+0.01%)	⬆️

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.