Adopt beartype runtime type checking across lcm

.github/workflows/main.yml

@@ -28,7 +28,7 @@ jobs:
       - uses: actions/checkout@v6
       - uses: prefix-dev/setup-pixi@v0.9.5
         with:
-          pixi-version: v0.67.2
+          pixi-version: v0.68.1
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           environments: tests-cpu
@@ -59,7 +59,7 @@ jobs:
       - uses: actions/checkout@v6
       - uses: prefix-dev/setup-pixi@v0.9.5
         with:
-          pixi-version: v0.67.2
+          pixi-version: v0.68.1
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           environments: type-checking
@@ -82,7 +82,7 @@ jobs:
       - uses: actions/checkout@v6
       - uses: prefix-dev/setup-pixi@v0.9.5
         with:
-          pixi-version: v0.67.2
+          pixi-version: v0.68.1
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           environments: tests-cuda12
@@ -101,7 +101,7 @@ jobs:
       - uses: actions/checkout@v6
       - uses: prefix-dev/setup-pixi@v0.9.5
         with:
-          pixi-version: v0.67.2
+          pixi-version: v0.68.1
           cache: true
           cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
           environments: tests-cuda12
@@ -116,7 +116,7 @@ jobs:
   #     - uses: actions/checkout@v6
   #     - uses: prefix-dev/setup-pixi@v0.9.5
   #       with:
-  #         pixi-version: v0.67.2
+  #         pixi-version: v0.68.1
   #         cache: true
   #         cache-write: ${{ github.event_name == 'push' && github.ref_name == 'main' }}
   #         environments: tests-cpu

AGENTS.md

@@ -275,13 +275,13 @@ loaded = SimulationResult.from_pickle("path/to/file.pkl")
 
 ### Initial Conditions Format
 
-Initial conditions use a flat dictionary with state names plus `"regime"`:
+Initial conditions use a flat dictionary with state names plus `"regime_id"`:
 
 ```python
 initial_conditions = {
     "wealth": jnp.array([1.0, 2.0, 3.0]),
     "health": jnp.array([0.5, 0.8, 0.3]),
-    "regime": jnp.array([RegimeId.working, RegimeId.working, RegimeId.retired]),
+    "regime_id": jnp.array([RegimeId.working, RegimeId.working, RegimeId.retired]),
 }
 ```
 
@@ -527,11 +527,16 @@ Code structure should be self-evident from function names and ordering.
 - **Helper function names follow `{verb}_{qualifier}_noun` patterns.** E.g.,
   `get_irreg_coordinate`, `find_irreg_coordinate`, `get_linspace_coordinate` — not
   `get_coordinate_irreg`.
-- **Use `@overload` when a function accepts both scalar and array inputs.** When a
-  function works with both `ScalarFloat` and `Array`, add overload declarations so the
-  type checker can track `(ScalarFloat) -> ScalarFloat` and `(Array) -> Array`
-  separately. Concrete subclass methods need their own overloads too (not just the
-  abstract base).
+- **Pick the single narrowest jaxtyping alias — never scalar/array `@overload` pairs,
+  never `ScalarX | XND` unions.** ty erases jaxtyping shape annotations: `ScalarFloat`,
+  `Float1D`, and `FloatND` all reveal as `Array`, so scalar/array `@overload` pairs and
+  `ScalarFloat | FloatND`-style unions add zero static precision — they are pure noise.
+  At runtime, beartype treats a 0-d float array as satisfying both `ScalarFloat` and
+  `FloatND`, so `ScalarFloat ⊆ FloatND` and the union is redundant. Annotate each slot
+  with the one alias that matches its genuine rank: `ScalarFloat`/`ScalarInt` for
+  fixed-0-d, `Float1D`/`Int1D` for fixed-1-d, `FloatND`/`IntND` for genuinely rank-
+  polymorphic. Never use a bare `Array` annotation — always reach for the narrowest
+  `lcm.typing` alias.
 - **`func` for callable abbreviations** — use `func`, `func_name`, `func_params` (never
   `fn`). Full word `function(s)` in dataclass field names and public method names.
 - **Singular `state_names` / `action_names`** — not `states_names` / `actions_names`.

benchmarks/_gpu_mem.py

@@ -26,6 +26,11 @@ class MahlerYumGpuPeakMem(GpuPeakMem):
 # Project root: the directory containing the benchmarks/ package.
 _PROJECT_ROOT = Path(__file__).resolve().parent.parent
 
+# Marks the peak-memory line on the subprocess's stdout. The subprocess imports
+# lcm, whose beartype claw can emit diagnostics to stdout, so the parent locates
+# this line instead of parsing stdout wholesale.
+_PEAK_MARKER = "__PEAK_BYTES_IN_USE__"
+
 
 def measure_gpu_peak(bench_module: str, bench_class: str) -> int:
     """Run a benchmark in a subprocess and return peak GPU bytes.
@@ -58,7 +63,15 @@ def measure_gpu_peak(bench_module: str, bench_class: str) -> int:
             f"stderr: {result.stderr!r}"
         )
         raise RuntimeError(msg)
-    return int(result.stdout.strip())
+    for line in result.stdout.splitlines():
+        if line.startswith(_PEAK_MARKER):
+            return int(line.removeprefix(_PEAK_MARKER).strip())
+    msg = (
+        "GPU memory subprocess produced no peak-bytes line.\n"
+        f"stdout: {result.stdout!r}\n"
+        f"stderr: {result.stderr!r}"
+    )
+    raise RuntimeError(msg)
 
 
 def _track_gpu_peak_mem(self):
@@ -104,4 +117,4 @@ def setup(self):
     import jax
 
     stats = jax.local_devices()[0].memory_stats()
-    print(stats["peak_bytes_in_use"])
+    print(f"{_PEAK_MARKER} {stats['peak_bytes_in_use']}")

benchmarks/bench_mahler_yum.py

@@ -29,7 +29,7 @@ def _build(self):
         self.model_params = {"alive": common_params}
         self.initial_conditions = {
             **initial_states,
-            "regime": jnp.full(
+            "regime_id": jnp.full(
                 _N_SUBJECTS,
                 self.model.regime_names_to_ids["alive"],
                 dtype=jnp.int32,

benchmarks/bench_precautionary_savings.py

@@ -33,7 +33,7 @@ def _make_initial_conditions(n_subjects):
         "age": jnp.full(n_subjects, 20.0),
         "wealth": jnp.full(n_subjects, 5.0),
         "income": jnp.full(n_subjects, 0.0),
-        "regime": jnp.zeros(n_subjects, dtype=jnp.int32),
+        "regime_id": jnp.zeros(n_subjects, dtype=jnp.int32),
     }
 
 

docs/development/benchmarking.md

@@ -129,7 +129,7 @@ class TimeMyModel:
         self.model_params = my_model.get_params()
         self.initial_conditions = {
             "wealth": jnp.full(1_000, 5.0),
-            "regime": jnp.zeros(1_000, dtype=jnp.int32),
+            "regime_id": jnp.zeros(1_000, dtype=jnp.int32),
         }
 
         # JIT warmup (timed separately)

docs/examples/mahler_yum_2024.md

@@ -40,7 +40,7 @@ result = MAHLER_YUM_MODEL.simulate(
     params={"alive": params},
     initial_conditions={
         **initial_states,
-        "regime": jnp.full(
+        "regime_id": jnp.full(
             n_subjects,
             MAHLER_YUM_MODEL.regime_names_to_ids["alive"],
         ),

docs/examples/mortality.md

@@ -26,7 +26,7 @@ result = model.simulate(
     initial_conditions={
         "age": jnp.full(100, model.ages.values[0]),
         "wealth": jnp.linspace(1, 100, 100),
-        "regime": jnp.full(100, model.regime_names_to_ids["working_life"]),
+        "regime_id": jnp.full(100, model.regime_names_to_ids["working_life"]),
     },
     period_to_regime_to_V_arr=None,
     seed=1234,

docs/examples/precautionary_savings.md

@@ -27,7 +27,7 @@ result = model.simulate(
         "age": jnp.full(100, model.ages.values[0]),
         "wealth": jnp.linspace(1, 10, 100),
         "income": jnp.zeros(100),
-        "regime": jnp.full(100, model.regime_names_to_ids["alive"]),
+        "regime_id": jnp.full(100, model.regime_names_to_ids["alive"]),
     },
     period_to_regime_to_V_arr=None,
 )

docs/examples/precautionary_savings_health.md

@@ -32,7 +32,7 @@ result = model.simulate(
         "age": jnp.full(1_000, model.ages.values[0]),
         "wealth": jnp.full(1_000, 1.0),
         "health": jnp.full(1_000, 1.0),
-        "regime": jnp.full(1_000, model.regime_names_to_ids["working_life"]),
+        "regime_id": jnp.full(1_000, model.regime_names_to_ids["working_life"]),
     },
     period_to_regime_to_V_arr=None,
 )

docs/examples/tiny.md

@@ -25,7 +25,7 @@ params = get_params()
 
 initial_df = pd.DataFrame(
     {
-        "regime": "working_life",
+        "regime_name": "working_life",
         "age": model.ages.values[0],
         "wealth": np.linspace(1, 20, 100),
     }

docs/user_guide/benchmarking.md

@@ -91,7 +91,7 @@ class TimeSolveSimulate:
         self.initial_conditions = {
             "age": jnp.full(500, 25.0),
             "wealth": jnp.full(500, 5.0),
-            "regime": jnp.zeros(500, dtype=jnp.int32),
+            "regime_id": jnp.zeros(500, dtype=jnp.int32),
         }
 
         # --- JAX warmup --------------------------------------------------

docs/user_guide/pandas_interop.md

@@ -12,12 +12,12 @@ is DataFrame in, DataFrame out.
 ## Initial Conditions as a DataFrame
 
 Pass a pandas DataFrame directly to `simulate()` as `initial_conditions`. One row per
-agent, one column per state variable, plus a `"regime"` column:
+agent, one column per state variable, plus a `"regime_name"` column:
 
 ```python
 df = pd.DataFrame(
     {
-        "regime": ["working", "working", "retired"],
+        "regime_name": ["working", "working", "retired"],
         "wealth": [10.0, 50.0, 30.0],
         "health": ["good", "bad", "good"],
         "age": [25.0, 25.0, 25.0],
@@ -31,7 +31,7 @@ result = model.simulate(
 )
 ```
 
-- `"regime"` column is required. Use regime names as strings (e.g., `"working"`).
+- `"regime_name"` column is required. Use regime names as strings (e.g., `"working"`).
 - Discrete states use string labels from the model's categorical classes (e.g., `"good"`
   instead of `0`). Labels are validated and mapped to integer codes automatically.
 - Continuous states pass through as-is.

docs/user_guide/solving_and_simulating.md

@@ -74,7 +74,7 @@ import pandas as pd
 
 df = pd.DataFrame(
     {
-        "regime": ["working_life", "working_life", "retirement", "working_life"],
+        "regime_name": ["working_life", "working_life", "retirement", "working_life"],
         "age": [25.0, 25.0, 25.0, 25.0],
         "wealth": [1.0, 5.0, 10.0, 20.0],
         "health": ["good", "bad", "bad", "good"],  # string labels, auto-converted
@@ -102,7 +102,7 @@ initial_conditions = {
     "age": jnp.array([25.0, 25.0, 25.0, 25.0]),
     "wealth": jnp.array([1.0, 5.0, 10.0, 20.0]),
     "health": jnp.array([0, 1, 1, 0]),  # integer codes for discrete states
-    "regime": jnp.array(
+    "regime_id": jnp.array(
         [
             RegimeId.working_life,
             RegimeId.working_life,
@@ -114,7 +114,7 @@ initial_conditions = {
 ```
 
 - Every non-shock state must have an entry.
-- `"regime"` must be included, with integer codes from the `regime_id_class`.
+- `"regime_id"` must be included, with integer codes from the `regime_id_class`.
 - All arrays must have the same length (= number of agents).
 - Shock states are drawn automatically.
 
@@ -140,7 +140,7 @@ Subjects can start at different ages:
 initial_conditions = {
     "age": jnp.array([40.0, 60.0]),
     "wealth": jnp.array([50.0, 50.0]),
-    "regime": jnp.array(
+    "regime_id": jnp.array(
         [
             model.regime_names_to_ids["working_life"],
             model.regime_names_to_ids["working_life"],
@@ -160,7 +160,7 @@ earlier periods are omitted, not filled with placeholders.
 df = result.to_dataframe()
 ```
 
-Returns a pandas DataFrame with columns: `subject_id`, `period`, `age`, `regime`,
+Returns a pandas DataFrame with columns: `subject_id`, `period`, `age`, `regime_name`,
 `value`, plus all states and actions. Discrete variables are pandas Categorical with
 string labels.
 
@@ -240,7 +240,7 @@ params = {
 
 # 3. Prepare initial conditions as a DataFrame
 initial_df = pd.DataFrame({
-    "regime": "working_life",
+    "regime_name": "working_life",
     "age": model.ages.values[0],
     "wealth": np.linspace(1, 50, 100),
 })