Skip to content

Fix issues with RNGs in Scan#2039

Open
ricardoV94 wants to merge 2 commits intopymc-devs:v3from
ricardoV94:rng_scan_bugfixes
Open

Fix issues with RNGs in Scan#2039
ricardoV94 wants to merge 2 commits intopymc-devs:v3from
ricardoV94:rng_scan_bugfixes

Conversation

@ricardoV94
Copy link
Copy Markdown
Member

Fixes two unrelated issues found in pymc-devs/pymc#8243

  1. Our rewrite sidestep_unused_rng_consumer could be too eager. Sometimes the output and the rng update are present in separate nodes that will eventually be merged together. Scan was doing this by cloning it separately due to implicit update machinery (deprecated API) that TruncatedRV in PyMC was still using. The explicit API doesn't clone/duplicate nodes, but we should still be on the guard.

The fix is to not perform the rewrite if the input RNG has multiple clients. This is either a likely wrong graph, or something that will be merged and can then be rewritten or not.

  1. scan (helper function) with constant n_steps=1, eagerly returns an "unrolled" scan. This was not handling untraced_sit_sot (the proper API for RNG updates) and was passing back internal dummy inputs, instead of the outer ones. This was revealed when we switched TruncatedRV to use the new API.

Could destroy update information in duplicated graphs (like Scan likes
to produce)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant