Model.n_subjects (review-only mirror of #340 with aca-model held flat)

benchmarks/bench_aca_baseline.py

@@ -54,7 +54,7 @@ def _build() -> tuple[object, object, object]:
         get_benchmark_params,
     )
 
-    model = create_benchmark_model()
+    model = create_benchmark_model(n_subjects=_N_SUBJECTS)
     _, model_params = get_benchmark_params(model=model)
     initial_conditions = get_benchmark_initial_conditions(
         model=model, n_subjects=_N_SUBJECTS, seed=0

pyproject.toml

@@ -98,7 +98,7 @@ tests-cuda13 = { features = [ "tests", "cuda13" ], solve-group = "cuda13" }
 tests-metal = { features = [ "tests", "metal" ], solve-group = "metal" }
 type-checking = { features = [ "type-checking", "tests" ], solve-group = "default" }
 [tool.pixi.feature.benchmarks.pypi-dependencies]
-aca-model = { git = "https://github.com/OpenSourceEconomics/aca-model.git", rev = "134286108b7445f3e17e8824bcdd1739a98b6089" }
+aca-model = { git = "https://github.com/OpenSourceEconomics/aca-model.git", rev = "3f215a2c44237b9fa3fa74bf78ef93c8d695a517" }
 [tool.pixi.feature.cuda12]
 platforms = [ "linux-64" ]
 system-requirements = { cuda = "12" }

src/lcm/grids/discrete.py

@@ -48,5 +48,15 @@ def batch_size(self) -> int:
         return self.__batch_size
 
     def to_jax(self) -> Int1D:
-        """Convert the grid to a Jax array."""
-        return jnp.array(self.codes)
+        """Convert the grid to a Jax array.
+
+        Discrete state/action codes are pinned to `int32` regardless of the
+        ambient `jax_enable_x64` setting. `jnp.array([...])` would otherwise
+        produce `int32` in 32-bit mode and `int64` in x64 mode, and
+        downstream values (transitions, V-array indexing, action lookups)
+        inherit that ambiguity — which silently splits the JIT cache into
+        per-period int32/int64 variants and breaks any AOT-compiled
+        program that ships a single signature. `int32` covers any realistic
+        category count and matches the `MISSING_CAT_CODE` sentinel.
+        """
+        return jnp.array(self.codes, dtype=jnp.int32)

src/lcm/model.py

@@ -1,6 +1,7 @@
 """Collection of classes that are used by the user to define the model and grids."""
 
 import dataclasses
+import logging
 from collections.abc import Mapping
 from pathlib import Path
 from types import MappingProxyType
@@ -30,6 +31,7 @@
 )
 from lcm.regime import Regime
 from lcm.regime_building.processing import InternalRegime
+from lcm.simulation.compile import compile_all_simulate_functions
 from lcm.simulation.initial_conditions import validate_initial_conditions
 from lcm.simulation.result import SimulationResult, get_simulation_output_dtypes
 from lcm.simulation.simulate import simulate
@@ -85,6 +87,16 @@ class Model:
     fixed_params: UserParams
     """Parameters fixed at model initialization."""
 
+    n_subjects: int | None = None
+    """Expected simulate batch size; enables AOT compile of simulate functions.
+
+    When set, the first matching `simulate(...)` call AOT-compiles all simulate
+    functions for batch shape `n_subjects` in parallel. Subsequent calls with the
+    same size reuse the compiled programs. Calls with a mismatching size warn
+    once per size and fall back to the runtime-traced path. `None` keeps the
+    purely lazy behaviour.
+    """
+
     _params_template: ParamsTemplate
     """Template for the model parameters."""
 
@@ -100,6 +112,7 @@ def __init__(
         derived_categoricals: Mapping[FunctionName, DiscreteGrid] = MappingProxyType(
             {}
         ),
+        n_subjects: int | None = None,
     ) -> None:
         """Initialize the Model.
 
@@ -115,17 +128,27 @@ def __init__(
                 not in states/actions. Broadcast to all regimes (merged with
                 each regime's own `derived_categoricals`). Raises if a regime
                 already has a conflicting entry.
+            n_subjects: Expected simulate batch size; if set, the first matching
+                `simulate(...)` call AOT-compiles all simulate functions for
+                batch shape `n_subjects` in parallel. `None` keeps the purely
+                lazy behaviour.
 
         """
         self.description = description
         self.ages = ages
         self.n_periods = ages.n_periods
         self.fixed_params = ensure_containers_are_immutable(fixed_params)
+        self.n_subjects = n_subjects
+        self._simulate_compile_cache: dict[
+            int, MappingProxyType[RegimeName, InternalRegime]
+        ] = {}
+        self._warned_n_subjects: set[int] = set()
 
         validate_model_inputs(
             n_periods=self.n_periods,
             regimes=regimes,
             regime_id_class=regime_id_class,
+            n_subjects=n_subjects,
         )
         self.regime_names_to_ids = MappingProxyType(
             dict(
@@ -240,6 +263,46 @@ def solve(
             )
         return period_to_regime_to_V_arr
 
+    def _resolve_simulate_internal_regimes(
+        self,
+        *,
+        actual_n_subjects: int,
+        internal_params: InternalParams,
+        log: logging.Logger,
+        max_compilation_workers: int | None,
+    ) -> MappingProxyType[RegimeName, InternalRegime]:
+        """Return internal_regimes to use for simulate; AOT cache when matching.
+
+        Returns the original `internal_regimes` when `n_subjects` is `None` or
+        when the actual batch size mismatches the declared one (logging a
+        warning once per mismatching size). Otherwise builds and caches the
+        AOT-compiled regimes for the matching size.
+        """
+        if self.n_subjects is None:
+            return self.internal_regimes
+        if actual_n_subjects != self.n_subjects:
+            if actual_n_subjects not in self._warned_n_subjects:
+                log.warning(
+                    "simulate called with n_subjects=%d but model declared "
+                    "n_subjects=%d; falling back to runtime compile.",
+                    actual_n_subjects,
+                    self.n_subjects,
+                )
+                self._warned_n_subjects.add(actual_n_subjects)
+            return self.internal_regimes
+        if self.n_subjects not in self._simulate_compile_cache:
+            self._simulate_compile_cache[self.n_subjects] = (
+                compile_all_simulate_functions(
+                    internal_regimes=self.internal_regimes,
+                    internal_params=internal_params,
+                    ages=self.ages,
+                    n_subjects=self.n_subjects,
+                    max_compilation_workers=max_compilation_workers,
+                    logger=log,
+                )
+            )
+        return self._simulate_compile_cache[self.n_subjects]
+
     def simulate(
         self,
         *,
@@ -341,10 +404,17 @@ def simulate(
                     )
                     exc.add_note(f"Snapshot saved to {snap_dir}")
                 raise
+        actual_n_subjects = len(next(iter(initial_conditions.values())))
+        simulate_internal_regimes = self._resolve_simulate_internal_regimes(
+            actual_n_subjects=actual_n_subjects,
+            internal_params=internal_params,
+            log=log,
+            max_compilation_workers=max_compilation_workers,
+        )
         result = simulate(
             internal_params=internal_params,
             initial_conditions=initial_conditions,
-            internal_regimes=self.internal_regimes,
+            internal_regimes=simulate_internal_regimes,
             regime_names_to_ids=self.regime_names_to_ids,
             logger=log,
             period_to_regime_to_V_arr=period_to_regime_to_V_arr,

src/lcm/model_processing.py

@@ -147,8 +147,11 @@ def validate_model_inputs(
     n_periods: int,
     regimes: Mapping[RegimeName, Regime],
     regime_id_class: type,
+    n_subjects: int | None = None,
 ) -> None:
     """Validate model constructor inputs."""
+    _fail_if_invalid_n_subjects(n_subjects=n_subjects)
+
     # Early exit if regimes are not lcm.Regime instances
     if not all(isinstance(regime, Regime) for regime in regimes.values()):
         raise ModelInitializationError(
@@ -201,6 +204,19 @@ def validate_model_inputs(
         raise ModelInitializationError(msg)
 
 
+def _fail_if_invalid_n_subjects(*, n_subjects: int | None) -> None:
+    """Raise TypeError if non-int, ValueError if non-positive."""
+    if n_subjects is None:
+        return
+    # `bool` is a subclass of `int`; reject explicitly so True/False don't slip through.
+    if not isinstance(n_subjects, int) or isinstance(n_subjects, bool):
+        msg = f"n_subjects must be an int or None, got {type(n_subjects).__name__}."
+        raise TypeError(msg)
+    if n_subjects <= 0:
+        msg = f"n_subjects must be a positive integer, got {n_subjects}."
+        raise ValueError(msg)
+
+
 def _validate_all_variables_used(regimes: Mapping[RegimeName, Regime]) -> list[str]:
     """Validate that all states and actions are used somewhere in each regime.
 

src/lcm/pandas_utils.py

@@ -153,7 +153,8 @@ def initial_conditions_from_dataframe(  # noqa: C901
         for col, arr in result_arrays.items()
     }
     initial_conditions["regime"] = jnp.array(
-        df["regime"].map(dict(regime_names_to_ids)).to_numpy()
+        df["regime"].map(dict(regime_names_to_ids)).to_numpy(),
+        dtype=jnp.int32,
     )
 
     return initial_conditions

src/lcm/regime_building/argmax.py

@@ -43,7 +43,7 @@ def argmax_and_max(
         # When there are no dimensions to reduce over, return:
         # - index 0 (trivial argmax since there's only one element)
         # - the array itself (already the maximum)
-        return jnp.array(0), a
+        return jnp.array(0, dtype=jnp.int32), a
 
     # Move axis over which to compute the argmax to the back and flatten last dims
     # ==================================================================================
@@ -65,7 +65,7 @@ def argmax_and_max(
     max_value_mask = a == _max
     if where is not None:
         max_value_mask = jnp.logical_and(max_value_mask, where)
-    _argmax = jnp.argmax(max_value_mask, axis=-1)
+    _argmax = jnp.argmax(max_value_mask, axis=-1).astype(jnp.int32)
 
     return _argmax, _max.reshape(_argmax.shape)