Skip to content

Handle compute_map=None in Scan #1435

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jun 2, 2025

Conversation

zaxtax
Copy link
Contributor

@zaxtax zaxtax commented May 31, 2025

Description

This should also fix some regressions in the recent release. Add a None check for compute_map in make_thunk

Related Issue

Closes #1425

Type of change

  • New feature / enhancement
  • Bug fix
  • Documentation
  • Maintenance
  • Other (please specify):

@zaxtax zaxtax requested a review from ricardoV94 May 31, 2025 15:21
Copy link

codecov bot commented May 31, 2025

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 82.12%. Comparing base (92eef5e) to head (914e6e4).
Report is 1 commits behind head on main.

Additional details and impacted files

Impacted file tree graph

@@           Coverage Diff           @@
##             main    #1435   +/-   ##
=======================================
  Coverage   82.12%   82.12%           
=======================================
  Files         211      211           
  Lines       49751    49752    +1     
  Branches     8817     8818    +1     
=======================================
+ Hits        40857    40859    +2     
  Misses       6714     6714           
+ Partials     2180     2179    -1     
Files with missing lines Coverage Δ
pytensor/scan/op.py 84.83% <100.00%> (+0.01%) ⬆️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@@ -1500,6 +1500,10 @@ def make_thunk(self, node, storage_map, compute_map, no_recycling, impl=None):
node_input_storage = [storage_map[r] for r in node.inputs]
node_output_storage = [storage_map[r] for r in node.outputs]

# HACK: Here to handle Blockwise Scans
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's not a hack, compute map for inner graphs is useless.

Instead of creating a compute_map make an if/else below where rval is returned and return an rval that doesn't try to assign to compute_map if it was None

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ricardoV94 is this change alright?

@ricardoV94 ricardoV94 changed the title Adding workaround so more graphs can be vectorized Handle compute_map=None in Scan Jun 2, 2025
@ricardoV94 ricardoV94 added the bug Something isn't working label Jun 2, 2025
@ricardoV94 ricardoV94 merged commit 864ebd1 into pymc-devs:main Jun 2, 2025
73 checks passed
@zaxtax zaxtax deleted the blockwise_scan_hack branch June 2, 2025 10:39
@ricardoV94
Copy link
Member

Thanks @zaxtax

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
bug Something isn't working
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Vectorize fails on scan
2 participants