Address review: docstrings, error-message fix, boundary tests

- `_transition_checks.py` module docstring no longer references a
  non-existent `regime_building/static_checks.py`; it describes the
  actual distinction (runtime numerical checks vs. construction-time
  regime-spec validators).
- `regime_building/transitions.py`: the `collect_state_transitions`
  docstring said ShockGrid states get a `lambda: None` stub; the code
  skips them entirely. Docstring corrected.
- `solution/validate_V.py`: `Dict` → `dict` in a Returns section to
  match the `dict[str, Any]` annotation.
- `_transition_checks.py` `_validate_no_reachable_incomplete_targets`:
  drop the override that, when a target regime was absent from the
  source's `state_transitions`, listed every state of the target as
  "missing" — including non-stochastic states that need no explicit
  entry. The preceding line already computes the correct stochastic-
  only set for that case.
- `user_regime.py`: add a scope-boundary note to
  `_validate_function_output_grid_indexing` — the AST check is
  deliberately best-effort and should be deleted rather than hardened
  if it ever produces false positives.
- Replace the two contentless happy-path validator tests with boundary
  inputs: values at the inclusive [0, 1] bounds and row sums just
  inside the sum-to-1 tolerance, so "does not raise" pins the
  tolerance/bound logic instead of being a bare smoke check.

Author: hmgaudecker

src/lcm/_transition_checks.py

@@ -7,9 +7,10 @@
 regimes, and no positive probability to a target with incomplete stochastic
 transitions.
 
-Distinct from `regime_building/static_checks.py`, which fires at canonical-
-regime construction time on Python source (AST). The checks here need a
-fully-built `Regime` plus user `flat_params` to run.
+These are runtime checks: they need a fully-built `Regime` plus user
+`flat_params` and evaluate the transition functions numerically. The
+construction-time regime-spec validators (`Regime.__post_init__`, which
+inspect grids, signatures, and Python source) are a separate concern.
 
 """
 
@@ -332,8 +333,6 @@ def _validate_no_reachable_incomplete_targets(
         if not jnp.any(regime_transition_probs[target_regime_name] > 0):
             continue
         missing = sorted(needs - set(transitions.get(target_regime_name, {})))
-        if target_regime_name not in transitions:
-            missing = sorted(f"next_{s}" for s in target_regime.variables.state_names)
         raise InvalidRegimeTransitionProbabilitiesError(
             f"Regime '{regime_name}' at age {age} has positive transition "
             f"probability to '{target_regime_name}', but '{regime_name}' "

src/lcm/regime_building/transitions.py

@@ -29,7 +29,8 @@ def collect_state_transitions(
     """Collect state transition functions from `state_transitions`.
 
     For each state, produces entries keyed as `f"next_{name}"`:
-    - ShockGrid -> stub `lambda: None`
+    - ShockGrid -> skipped (shock transitions are built directly in
+      `_process_regime_core`)
     - `None` -> auto-generated identity transition
     - Callable -> used directly
     - `MarkovTransition` -> used directly (callable via `__call__`)

src/lcm/solution/validate_V.py

@@ -202,7 +202,7 @@ def _summarize_diagnostics(
         age: Age for the summary header.
 
     Returns:
-        Dict with per-metric `"overall"` and `"by_dim"` entries plus a
+        dict with per-metric `"overall"` and `"by_dim"` entries plus a
         `"regime_probs"` mapping, suitable for `_format_diagnostic_summary`.
 
     """

src/lcm/user_regime.py

@@ -399,6 +399,12 @@ def _validate_function_output_grid_indexing(regime: Regime) -> list[str]:
     consumer is indexing a 0-d array by a scalar integer, which raises
     `IndexError` at trace time. The fix is to drop the redundant `[g]`
     in the consumer (or refactor `f` not to take `g`).
+
+    This check is deliberately best-effort: it catches the common
+    `func_output[discrete_grid]` subscript form and nothing else. It is not
+    meant to grow into a general correctness checker for user functions —
+    if it ever produces false positives, prefer deleting it over hardening
+    it to chase every way the pattern can hide.
     """
     function_output_names = set(regime.functions)
     discrete_grid_names = (

tests/test_regime.py

@@ -439,9 +439,21 @@ def _make_partner_probs_array():
     )
 
 
-def test_validate_transition_probs_valid():
+def test_validate_transition_probs_accepts_boundary_inputs():
+    """Inclusive [0, 1] bounds and row sums within the 1e-6 tolerance pass.
+
+    The first row is exactly `[0.0, 1.0]` — values at the inclusive bounds.
+    The last row sums to `1 - 5e-7`, just inside the `atol=1e-6` row-sum
+    tolerance. The validator must accept both without raising.
+    """
     model = get_stochastic_model(3)
-    arr = _make_partner_probs_array()
+    arr = jnp.array(
+        [
+            [[[0.0, 1.0], [1.0, 0.0]], [[0.3, 0.7], [0.6, 0.4]]],
+            [[[0.4, 0.6], [0.8, 0.2]], [[0.2, 0.8], [0.7, 0.3]]],
+            [[[0.5, 0.5], [0.9, 0.1]], [[0.3, 0.7], [0.5, 0.4999995]]],
+        ]
+    )
     validate_transition_probs(
         probs=arr, model=model, regime_name="working_life", state_name="partner"
     )

tests/test_validate_regime_transition_probs.py

@@ -22,12 +22,17 @@
 from lcm_examples.mortality import get_model, get_params
 
 
-def test_valid_probs_all_active():
-    """Valid probabilities with all regimes active pass validation."""
+def test_valid_probs_accept_boundary_inputs():
+    """Inclusive [0, 1] bounds and a sum within tolerance pass validation.
+
+    Subject 0 splits probability exactly `[1.0, 0.0]` — values at the
+    inclusive bounds. Subject 1 sums to `1 - 2.5e-6`, just inside the
+    `jnp.allclose` default tolerance. The validator must accept both.
+    """
     probs = MappingProxyType(
         {
-            "working_life": jnp.array([0.7, 0.6]),
-            "retirement": jnp.array([0.3, 0.4]),
+            "working_life": jnp.array([1.0, 0.4999975]),
+            "retirement": jnp.array([0.0, 0.5]),
         }
     )
     _validate_regime_transition_probs(