Keep the remaining caching in FrozenModel - #8385
Open
velochy wants to merge 4 commits into
Open
Conversation
`pymc.util.hashable` and `HashableWrapper` back the memoization helpers, and mishandle three kinds of values: - Looking up a key compared the wrapped objects, which raises `ValueError: truth value of an array is ambiguous` for arrays and for the containers holding them. Keys are now compared by type and by the hash `hashable` computes, so a cache can be keyed on a dict of arrays. - Sets fell through to being pickled whole, which is not stable across calls, so equal sets got different hashes and silently defeated the cache. They are now hashed by their elements, like lists and dicts. - `In`/`Out` wrap a variable with compilation options but are hashed by identity, and callers rebuild them on every call, so a key holding one never matched. They are now hashed by what they hold. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Repeated `pm.sample` on a frozen model still recompiled the initial-point and trace functions, so the caching a frozen model advertises never fired on the path that matters most. Both call sites now go through the model, which the cache-key fixes make possible. Repeated sampling of a frozen model now compiles nothing. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
`logp_dlogp_function` and `initial_point` were still built around private methods in the base model that only existed to give the frozen model's cache something to wrap, the same shape that `compile_fn` was cleaned up from. Declare both in the base model and implement them in each final instead. `Model` builds its function every call, `FrozenModel` builds it through a cached method it owns. The part worth sharing, assembling the `ValueGradFunction`, moves to a module level helper both call, so the duplication is the argument handling rather than the logic. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
6 tasks
Documentation build overview
7 files changed ·
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #8385 +/- ##
=======================================
Coverage 91.84% 91.84%
=======================================
Files 128 128
Lines 21227 21253 +26
=======================================
+ Hits 19495 19520 +25
- Misses 1732 1733 +1
🚀 New features to boost your workflow:
|
Both finals override all three, and the base model cannot be instantiated, so the stubs are unreachable by construction and only show up as uncovered lines. Same marker the stand-in ops in `logprob` use. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
Author
|
(comment drafted by AI)
The Why did the abstract methods never get any coverage? They kept telling everyone "not implemented" and their subclasses kept overriding them anyway. Classic middle-management. 📋 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Follow-up to the review on #8330, which asked for these two to get the same treatment
compile_fnalready got.logp_dlogp_functionandinitial_pointwere still built around private methods onBaseModelthat existed only to giveFrozenModel's cache something to wrap, so every model carried a seam for a feature only frozen models use.Both are now declared in
BaseModeland implemented in each final:Modelbuilds its function on every call, as it did before any of this.FrozenModelbuilds it through a cached method it owns (_value_grad_function,_initial_point_fn).The part worth sharing — assembling the
ValueGradFunction— moves to a module level_make_value_grad_functionthat both call, so what is duplicated is the argument handling rather than the logic.make_initial_point_fns_per_chainfollows the rename from_make_initial_pointto_initial_point_fn, which both finals define, so sampling reuses the compiled initial-point function on a frozen model exactly as before.Related Issue
Checklist
Type of change