Fold IID utility shocks out of the stored value (quadrature-average at solve) - #402
Closed
hmgaudecker wants to merge 1 commit into
Closed
Fold IID utility shocks out of the stored value (quadrature-average at solve)#402hmgaudecker wants to merge 1 commit into
hmgaudecker wants to merge 1 commit into
Conversation
…erage at solve) An IID process state entering only the current period's utility can now be declared `fold=True`, integrating it out at solve time (quadrature-averaged into the stored value) instead of keeping its discretization nodes as an axis of every stored V. This kills the multiplicative node-memory blow-up for models with several such shocks. Default `fold=False` is byte-identical to before. The reduction lives in the grid-search max-Q-over-a kernel: a folded shock is still an ordinary inner productmap axis THROUGH the max-over-actions / collective readout (every node evaluated exactly as today), then its axis is weighted-averaged away with the process's own quadrature weights before the result is returned — so the stored V loses the axis, exactly (same quadrature, reduced one step earlier). The collective divorce flag D is folded by jnp.any to keep its boolean contract. Validation rejects fold=True when the shock is read by a same-period gate / value-constraint / reference projection (E2/E3'), conditioned on by a next-period transition, combined with taste_shocks, on a non-GridSearch solver, or has runtime-supplied distribution params; a cross-regime check rejects a folded state that structurally persists into a reachable regime. A persistent (non-IID) process has no fold field, so fold=True on one is a type error. NOTE: the `ty` type-check hook cannot run in a git worktree (invalid environment.python), so this was committed with --no-verify; ty must be re-checked in the main tree post-merge. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Benchmark comparison (main → HEAD)Comparing
|
Member
Author
|
Superseded by #409, which targets |
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.
Fold IID utility shocks out of the stored value (quadrature-average at solve)
Draft — stacked on
collective-regimes-draft(the collective-regimes extension). Review only the single commit on this branch; the base carries the extension.Motivation
An IID shock that enters the current period's utility (a wage draw, match quality, a pregnancy/taste innovation) is today materialized as a stored grid axis of every value array. With several such shocks the stored
Vpicks up a multiplicative node factor (×27–81 in Eckstein–Keane–Lifshitz-scale collective models), which is what makes full-resolution solves intractable. But such a shock is only ever integrated out against its own quadrature weights — it need not be stored at all.What this PR does
Adds
fold: bool = Falseto the IID process declaration.fold=Truemeans "integrate this shock into the value at solve time; do not keep it as a stored axis." The shock is materialized only through the max-over-actions / collective readout, then averaged away by its own quadrature weights before the value is stored:This is exact, not an approximation — the same quadrature a stored shock would use, reduced one step earlier. The stored
Vloses the folded axis.Not the same as
ExtremeValueTasteShocks(#375): those fold a Gumbel shock on the discrete-action axis via a closed-formlogsumexp. This folds a Normal/LogNormal shock on a state-node axis via numericaljnp.average(weights=). Same concept (integrate out, don't store), different operator and different reduction axis — no shared code.Where the reduction lands (the load-bearing bit)
In
max_Q_over_a.py::_wrap_with_fold_reduction: the fold axes are carried like co-map axes through the innerproductmap/readout unchanged, thenjnp.average'd out after the max/collective-readout and before any co-mapvmapwrapping. Reduced from highest inner-position down (removing one never shifts a not-yet-reduced one). The collective stakeholder axis is trailing onVonly, so it never collides with a folded state axis; the divorce flagDis folded withjnp.anyto stay strictly boolean.backward_induction._get_regime_V_shapes_and_shardingsdrops folded axes from the stored-V shape/sharding; the continuation read is unchanged (the stored V already lacks the axis).Validation (rejects unfoldable declarations)
fold=Trueis rejected with a clear error when the shock is: not IID (AR(1) must be stored); read by a same-period gate / value-constraint / reference projection (E2/E3′ need unfolded per-node values — walks the gate/constraint DAG ancestors); conditioned on by any next-period transition (a folded shock is integrated out, nothing downstream may depend on its realization — walks the transition DAGs); combined withtaste_shocks; on a non-GridSearch solver; or carrying runtime-supplied distribution params. A folded state that structurally persists into a reachable regime raisesModelInitializationError.Tests
tests/regime_building/test_fold_iid_shocks.py(10, written first):fold=Falsestored axis vsfold=True); assert value-identical (assert_allclose) with the foldedVcarrying exactly one fewer axis (single shock(5,)→(); 3-shock model drops 3 axes).assert_array_equal).Verification
tests/regime_building109 passed.pytest tests/solution tests/simulation: 570 passed, 5 skipped, 2 xfailed — baseline exactly (default path byte-identical; everything gates onfold=True).ty(run in the main tree — the hook can't run in a git worktree): All checks passed (one real annotation fix included:solvers.pycasts the folded grid to_ContinuousStochasticProcessbefore reading its transition probs).🤖 Generated with Claude Code