Age specialization: per-age regime functions/constraints and age-dependent continuous-state grids by hmgaudecker · Pull Request #398 · OpenSourceEconomics/pylcm

hmgaudecker · 2026-07-06T16:16:45Z

Summary

Adds age specialization to pylcm: one regime definition whose functions,
constraints, or continuous-state grids vary with age. Two user-facing markers
share a common base (_AgeSpecialized), each wrapping a build(age) factory
plus a signature(age) function — cells with equal signatures share one built
object, so compiled programs are deduplicated across ages.

`AgeSpecializedFunction` — per-age functions and constraints

Swaps the implementation of a regime function or constraint by age. The
motivating use case is calendar-policy tracking in cohort models: a GETTSIM
tax function whose policy environment moves with the agent's age, rebuilt per
(cohort, age) cell without duplicating regimes.
Solution: Q_and_F resolves the specialized function per age during backward
induction; representative-age resolution in _process_regime_core keeps
templates and the params template consistent (including
Phased(AgeSpecializedFunction)).
Simulation: next_state is period-indexed so per-age transitions simulate
correctly; to_dataframe(additional_targets=...) rejects age-specialized
functions it cannot represent at a single age.
Guardrails (each rejected at model construction with a targeted error):
deterministic state transitions that are age-specialized (bare or via
MarkovTransition(...), including per-target and Phased forms); regime
transitions that are age-specialized, wrapped, or whose dependency graph
reads an age-specialized helper; terminal-regime markers.
Docs: user-guide page docs/user_guide/age_specialized.md.

`AgeSpecializedGrid` — age-dependent continuous-state grids

A continuous state's grid bounds may move with age while n_points (and
the grid class) stay fixed — shape-invariant, so V-array shapes, the rolling
continuation template, and kernel-dedup topology are untouched. The
motivating use case is an asset state with an age-dependent borrowing floor
a̲(age) (Guvenen & Smith 2014, and any life-cycle model with an age-varying
constraint): on a single fixed grid, cells below a tighter age's floor are
infeasible, and their -inf continuation poisons the value function through
interpolation (0 · (−inf) = NaN, propagated backward).
Solution: the current-period state axis is runtime data — each period's node
values are substituted into the shared compiled kernel without retracing;
the continuation V_{t+1} is interpolated on the target regime's
period-(t+1) grid, with the Q_and_F group signature folding in the
continuation-grid identity so periods never false-share a program.
Simulation: the forward pass interpolates the continuation on the same
per-period grids as the solve.
Validation: build(age) must return the same grid class and n_points at
every age (RegimeInitializationError otherwise); markers are accepted only
as state grids.
Design notes: DESIGN-age-dependent-grids.md.

Hardening

The function-specialization surface went through several rounds of external
adversarial review (ChatGPT-Pro code audit) driven by the downstream rmsm
project; every hole found (transitive regime-transition reads,
MarkovTransition-wrapped markers, params-template blind spots,
representative-age leaks into additional_targets) is closed with a
regression test. Downstream, an annual lifecycle model with per-(cohort, age)
GETTSIM policy environments solves and simulates on this branch, and its
estimation runs on GPU. The grid surface is exercised end-to-end by the
Guvenen–Smith (2014) replication in lcm-replications, where the age-tracking
asset floor eliminates the -inf/NaN poisoning of the fixed-grid solve.

Test plan

Full suite: 1229 passed, 48 skipped (pixi run -e tests-cpu tests);
ruff check and ty both fully clean.
Functions: tests/regime_building/test_age_specialization.py
(validation surface, sharing by signature),
tests/solution/test_age_specialized_solve.py (solve/simulate
equivalence against a hand-built per-age model).
Grids: tests/solution/test_age_specialized_grid_solve.py — an
age-invariant AgeSpecializedGrid reproduces the plain fixed-grid solve
bit-for-bit; an age-tracking floor solves with no NaN poisoning;
V monotone in wealth; simulation yields finite, positive consumption;
non-shape-invariant grids are rejected.
Equivalence gates: constant AgeSpecializedFunction signature
reproduces the unspecialized solution exactly; models without any
age-varying grid build byte-identically to the pre-feature path.

🤖 Generated with Claude Code

https://claude.ai/code/session_016BLLuvQwD2bdXgaTcJWWUm

…river) Two tests on the 4-CPU distributed seam: - test_distributed_solve_matches_single_device_per_type: the F6 correctness contract — a sharded solve must equal the single-device solve per type slice. Green now; the co-map fix must preserve it (catches a wrong device-local index). - test_distributed_solve_kernel_does_not_all_gather_continuation_v: the red driver — the backward-induction kernel currently all-gathers the continuation V across the type shard; it must read only its device-local slice. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

vmap_1d gains co_mapped_in_axes: an optional per-argument in_axes override so a pytree argument's leading axis can be mapped in lockstep with the mapped variables. The backward-induction co-map uses it to slice each next_regime_to_V_arr leaf to the device-local type, so the continuation-V interpolation reads only its own shard. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Fixed, distributed states (e.g. a permanent type sharded one block per device) never transition, so a regime's continuation value depends only on its own slice of the next-period V-array. The grid-search solve kernel now co-maps each such state with the matching axis of every next_regime_to_V_arr leaf that carries it: an outer vmap peels the leading axis off both the state grid and the continuation V, so the interpolation reads only the device-local slice and XLA inserts no all-gather of the full V-array onto every device. - max_Q_over_a splits the co-mapped (leading) states from the inner productmap and wraps it in per-state co-map vmaps; the V-interpolator drops those coordinates and Q_and_F omits the sliced next-states. - processing detects the co-mappable states (distributed and identity-transition) and builds per-state, per-leaf in_axes so a target regime that prunes the state keeps its full leaf. - Simulation is unchanged: it keeps the full continuation V (subjects are not type-aligned across devices), so the co-map applies to the solve path only. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…nuation-v

…seam (#385) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>