@categorical: require ScalarInt field annotations + runtime values by hmgaudecker · Pull Request #350 · OpenSourceEconomics/pylcm

hmgaudecker · 2026-05-11T12:41:17Z

Summary

Tightens @categorical so the annotation contract, runtime values, and validator all agree on ScalarInt. Fields annotated otherwise raise the new CategoricalDefinitionError at decoration time, and class- / instance-level attribute access (LaborSupply.work, LaborSupply().work) now return 0-d jnp.int32 scalars instead of Python ints.
Adds lcm.invert_regime_ids (in lcm.utils.containers) so the two pylcm dict-inversion sites can coerce ScalarInt keys to Python int without spelling out the comprehension. DiscreteGrid.codes likewise coerces at the boundary where hashable tuples are needed; the JAX-side surface (to_jax()) is unchanged.
Sweeps every @categorical-decorated class across src/lcm_examples/, tests/, and the user-guide notebooks/docs to ScalarInt. The Effort example in mahler_yum_2024/_model.py switches to a make_dataclass + type.__setattr__ shim so its 40 ScalarInt defaults bypass dataclass's mutable-default check.

Closes #349.

Test plan

pixi run -e tests-cpu tests — 957 passed, 5 skipped.
pixi run -e type-checking ty — clean.
prek run --all-files — clean.
Downstream cascade in aca-model verified locally (216 aca-model tests + 407-test workspace suite green with this pylcm branch as the submodule pointer); separate aca-model PR will follow.

🤖 Generated with Claude Code

… values (#349) The decorator used to read `__annotations__` only for field names and never validate them — every consumer wrote `field: int` even though codes flowed through JAX as `jnp.int32` everywhere downstream. The mismatch was a consistent lie ty couldn't catch, and downstream `jnp.int32(...)` wraps papered over the dtype gap at runtime. This change closes the gap at all three layers in lockstep: * **Annotation gate.** `@categorical` requires every field to be annotated `ScalarInt` (from `lcm.typing`); other annotations raise the new `CategoricalDefinitionError` at decoration time, naming the offending fields and pointing at the import. * **Runtime values.** Class- and instance-level attribute access return 0-d `jnp.int32` scalars. The decorator assigns `field(default=i, init=False)` (Python int placeholders for `dataclass(frozen=True)`'s mutable-default check) then overrides the class attributes with `jnp.int32(i)` via `type.__setattr__` post-decoration; `init=False` keeps instance `__dict__` empty so attribute lookup falls through to the class scalar. * **Validator.** `validate_category_class` now checks `isinstance(value, jax.Array) and value.shape == () and jnp.issubdtype(value.dtype, jnp.integer)`; the consecutiveness check coerces via `int(v)`. Hashability fix: JAX 0-d arrays aren't hashable, so the two pylcm dict-inversion sites (`simulation/simulate.py`, `initial_conditions.py`) and `DiscreteGrid.codes` now coerce to Python `int` at the boundary where hashable keys are needed. New helper `lcm.invert_regime_ids` lives in `lcm.utils.containers` and centralises the inversion pattern; downstream consumers can call it instead of hand-rolling `{int(v): k for k, v in ...}`. `RegimeNamesToIds` tightens to `MappingProxyType[RegimeName, ScalarInt]`. Sweep: every `@categorical`-decorated class in `src/lcm_examples`, `tests/`, `tests/test_models/`, and the user-guide notebooks/docs moves from `field: int` to `field: ScalarInt`. The `Effort` example in `mahler_yum_2024/_model.py` switches to the `make_dataclass + type.__setattr__` shim so its 40 `ScalarInt` defaults bypass `dataclass`'s mutable-default check. `tests/conftest.py` and `tests/test_grids.py` gain a `_make_dc` helper for hand-crafted test dataclasses with `ScalarInt` values. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

review-notebook-app · 2026-05-11T12:41:23Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

read-the-docs-community · 2026-05-11T12:44:03Z

Documentation build overview

📚 pylcm | 🛠️ Build #32648285 | 📁 Comparing 64d8659 against latest (99a5e31)

🔍 Preview build

32 files changed · ± 32 modified

± Modified

PR #350 tightens `@categorical` to require `ScalarInt` field annotations. The benchmarks-cuda12 env still pulls aca-model from the pre-#350 sha (`9ac2043`), which annotates fields as `int` and fails the new decoration-time gate during `setup_cache`. Point the pin at the aca-model #350 cascade branch (`feature/categorical-scalarint`, head `b807b28`) until both PRs land. After aca-model PR#10 and pylcm#350 merge, revert to a `main`-tracking sha. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

aca-model's `get_benchmark_params` returns three values (likely a moments tuple alongside the two pre-existing entries); the benchmark file was stuck on the older 2-tuple shape. aca-model's own `test_benchmark.py` runs against the up-to-date in-tree helper, so the drift only surfaces here once pylcm's benchmark CI pulls a fresh aca-model. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

github-actions · 2026-05-11T15:07:17Z

Benchmark comparison (main → HEAD)

Comparing 99a5e31d (main) → 64d8659a (HEAD)

Benchmark	Statistic	before	after	Ratio	Alert
aca-baseline	execution time	27.474 s	27.491 s	1.00
	peak GPU mem	509 MB	579 MB	1.14	❌
	compilation time	299.51 s	300.51 s	1.00
	peak CPU mem	7.65 GB	7.34 GB	0.96
Mahler-Yum	execution time	4.712 s	4.631 s	0.98
	peak GPU mem	529 MB	529 MB	1.00
	compilation time	14.59 s	14.73 s	1.01
	peak CPU mem	1.68 GB	1.68 GB	1.00
Precautionary Savings - Solve	execution time	50.8 ms	53.1 ms	1.05
	peak GPU mem	101 MB	101 MB	1.00
	compilation time	2.71 s	2.68 s	0.99
	peak CPU mem	1.13 GB	1.13 GB	1.00
Precautionary Savings - Simulate	execution time	126.7 ms	120.7 ms	0.95
	peak GPU mem	344 MB	344 MB	1.00
	compilation time	4.90 s	5.09 s	1.04
	peak CPU mem	1.31 GB	1.31 GB	1.00
Precautionary Savings - Solve & Simulate	execution time	145.2 ms	155.5 ms	1.07
	peak GPU mem	578 MB	578 MB	1.00
	compilation time	7.02 s	7.16 s	1.02
	peak CPU mem	1.28 GB	1.29 GB	1.01
Precautionary Savings - Solve & Simulate (irreg)	execution time	283.3 ms	283.6 ms	1.00
	peak GPU mem	2.19 GB	2.19 GB	1.00
	compilation time	7.58 s	7.52 s	0.99
	peak CPU mem	1.34 GB	1.35 GB	1.01

Class-attr access already documents the `Array(0, dtype=int32)` shape; the instance fall-through line repeats the same value without adding information about the decorator contract. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…_to_names Mirrors the forward direction `regime_names_to_ids` so the inverse is read symmetrically. No behavior change. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Matches the inverse direction's symmetry with `regime_names_to_ids` and the recent rename in `simulation.initial_conditions`. No behavior change. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

hmgaudecker

Autoreview.

- bench: simplify get_benchmark_params unpacking via index access. - categorical: tighten validate_category_class docstring to refer to `ScalarInt`s directly without restating the type. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

hmgaudecker · 2026-05-11T18:04:38Z

This completes #340.

Lots of changes (sorry, @mj023 @timmens), but almost all of them just the expected int -> ScalarInt, so just have a quick glance.

Just a workaround necessary for RegimeIds since they are mostly used in pure Python, but not worth an exception IMO.

timmens

Looks good to me!

Annotate the inverse-direction mapping with its own PEP-695 alias so both directions read symmetrically. Updates `log_regime_transitions` and the two `_collect_*_errors` helpers in `simulation.initial_conditions` to use it; the docstrings switch from "Mapping from..." to "Immutable mapping of..." per the MappingProxyType convention. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Cascades #350's squash merge (RegimeIdsToNames alias + renames). Conflicts in typing.py/logging.py/initial_conditions.py resolved in favor of main — feat/runtime-grid-extra-params had the older un-squashed history of the same content, so main's squash is the canonical form. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

hmgaudecker mentioned this pull request May 11, 2026

@categorical fields: int → ScalarInt (pylcm #350 cascade) OpenSourceEconomics/aca-model#10

Merged

3 tasks

hmgaudecker force-pushed the feat/categorical-scalarint branch from 4a05a82 to 6e187a7 Compare May 11, 2026 13:14

hmgaudecker mentioned this pull request May 11, 2026

IrregSpacedGrid: extra_param_names for user-side runtime scalars #351

Closed

4 tasks

hmgaudecker and others added 3 commits May 11, 2026 19:52

refactor(initial_conditions): rename ids_to_regime_names → regime_ids…

04c9533

…_to_names Mirrors the forward direction `regime_names_to_ids` so the inverse is read symmetrically. No behavior change. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

hmgaudecker commented May 11, 2026

View reviewed changes

Comment thread benchmarks/bench_aca_baseline.py Outdated

Comment thread src/lcm/grids/categorical.py Outdated

Comment thread src/lcm/grids/categorical.py Outdated

hmgaudecker requested review from mj023 and timmens and removed request for timmens May 11, 2026 18:02

timmens approved these changes May 12, 2026

View reviewed changes

Comment thread src/lcm/utils/logging.py Outdated

hmgaudecker merged commit 3eec679 into main May 12, 2026
10 checks passed

hmgaudecker deleted the feat/categorical-scalarint branch May 12, 2026 05:32

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

@categorical: require ScalarInt field annotations + runtime values#350

@categorical: require ScalarInt field annotations + runtime values#350
hmgaudecker merged 8 commits into
mainfrom
feat/categorical-scalarint

hmgaudecker commented May 11, 2026 •

edited

Loading

Uh oh!

review-notebook-app Bot commented May 11, 2026

Uh oh!

read-the-docs-community Bot commented May 11, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented May 11, 2026 •

edited

Loading

Uh oh!

hmgaudecker left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

hmgaudecker commented May 11, 2026

Uh oh!

timmens left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

hmgaudecker commented May 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

review-notebook-app Bot commented May 11, 2026

Uh oh!

read-the-docs-community Bot commented May 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Documentation build overview

Uh oh!

github-actions Bot commented May 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Benchmark comparison (main → HEAD)

Uh oh!

hmgaudecker left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

hmgaudecker commented May 11, 2026

Uh oh!

timmens left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

hmgaudecker commented May 11, 2026 •

edited

Loading

read-the-docs-community Bot commented May 11, 2026 •

edited

Loading

github-actions Bot commented May 11, 2026 •

edited

Loading