Extend GridConfig with distributed flags for non-pref_type discrete states

Add `lagged_labor_supply_distributed`, `claimed_ss_distributed`, and
`spousal_income_distributed` to `GridConfig` and thread them through
`build_states` into the inline-built `DiscreteGrid(...)` calls.
Enables the 2x2 (lagged_labor_supply x claimed_ss) and 3-way
(spousal_income) sharding configurations needed by the OOM /
performance experiment matrix on Marvin.

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

Author: hmgaudecker

src/aca_model/baseline/regimes/_common.py

@@ -409,18 +409,24 @@ def build_states(spec: RegimeSpec, grids: Grids) -> dict:
     states["hcc_persistent"] = grids.hcc_persistent
     states["hcc_transitory"] = grids.hcc_transitory
     states["spousal_income"] = DiscreteGrid(
-        SpousalIncome, batch_size=gc.n_spousal_income_batch_size
+        SpousalIncome,
+        batch_size=gc.n_spousal_income_batch_size,
+        distributed=gc.spousal_income_distributed,
     )
     states["pref_type"] = grids.pref_type
     if can_work:
         states["log_ft_wage_res"] = grids.wage_res
     if can_work and spec["his"] != "tied":
         states["lagged_labor_supply"] = DiscreteGrid(
-            LaggedLaborSupply, batch_size=gc.n_lagged_labor_supply_batch_size
+            LaggedLaborSupply,
+            batch_size=gc.n_lagged_labor_supply_batch_size,
+            distributed=gc.lagged_labor_supply_distributed,
         )
     if spec["ss"] == "choose":
         states["claimed_ss"] = DiscreteGrid(
-            ClaimedSS, batch_size=gc.n_claimed_ss_batch_size
+            ClaimedSS,
+            batch_size=gc.n_claimed_ss_batch_size,
+            distributed=gc.claimed_ss_distributed,
         )
     return states
 

src/aca_model/config.py

@@ -31,12 +31,19 @@ class GridConfig:
     # intermediate by 12x on hosts where the unsplayed kernel doesn't fit.
     n_assets_batch_size: int = 0
     n_aime_batch_size: int = 1
-    # Sharding flag for the `pref_type` discrete grid: pylcm distributes
-    # the grid across devices when `distributed=True`. Sharding is only
-    # supported on discrete state grids; continuous axes (`assets`,
-    # `aime`, `wage_res`, `hcc_*`) compile to an all-gather of the full
-    # V-array per device and are rejected at grid construction.
+    # Sharding flags for discrete state grids. pylcm distributes the
+    # grid across available devices when the flag is `True`. Sharding
+    # is only supported on discrete state grids; continuous axes
+    # (`assets`, `aime`, `wage_res`, `hcc_*`) compile to an all-gather
+    # of the full V-array per device and are rejected at grid
+    # construction. Mutually exclusive with `batch_size>0` on the same
+    # axis (pylcm rejects the combination). The non-`pref_type` flags
+    # route through `baseline/regimes/_common.py:build_states` to the
+    # inline-built `DiscreteGrid(...)` calls.
     pref_type_distributed: bool = False
+    lagged_labor_supply_distributed: bool = False
+    claimed_ss_distributed: bool = False
+    spousal_income_distributed: bool = False
     # `batch_size` on the inline-constructed discrete state grids —
     # health, spousal_income, lagged_labor_supply, claimed_ss. These
     # are read in `build_states` via `grids.grid_config`. Setting any

tests/test_model_creation.py

@@ -1,6 +1,7 @@
 """Tests for baseline model creation and regime structure."""
 
 from collections.abc import Mapping
+from dataclasses import replace
 
 import pytest
 from helpers.model import make_aca_model, make_baseline_model
@@ -274,3 +275,38 @@ def test_baseline_model_creates() -> None:
     """Baseline model creates successfully without PolicyVariant."""
     model = make_baseline_model(n_subjects=1)
     assert len(model.user_regimes) == 19
+
+
+@pytest.mark.parametrize(
+    ("config_field", "state_name"),
+    [
+        ("lagged_labor_supply_distributed", "lagged_labor_supply"),
+        ("claimed_ss_distributed", "claimed_ss"),
+        ("spousal_income_distributed", "spousal_income"),
+    ],
+)
+def test_discrete_state_distributed_flag_propagates_to_regime(
+    config_field: str, state_name: str
+) -> None:
+    """`GridConfig.<axis>_distributed=True` sets `distributed=True` on the
+    `DiscreteGrid` for that axis in every regime that carries it."""
+    gc = replace(BENCHMARK_GRID_CONFIG, **{config_field: True})
+    grids = build_grids(
+        grid_config=gc,
+        fixed_params=_FIXED_PARAMS,
+        wage_params=_WAGE_PARAMS,
+        pref_type_grid=DiscreteGrid(BenchmarkPrefType),
+    )
+    regime = _build_regime("retiree_dimc_choose_canwork", grids)
+    assert regime.states[state_name].distributed is True
+
+
+@pytest.mark.parametrize(
+    "state_name",
+    ["lagged_labor_supply", "claimed_ss", "spousal_income"],
+)
+def test_discrete_state_distributed_flag_defaults_to_false(state_name: str) -> None:
+    """`distributed` on inline-built discrete states defaults to `False` so
+    configurations that do not opt in see no behaviour change."""
+    regime = build_regime("retiree_dimc_choose_canwork")
+    assert regime.states[state_name].distributed is False