Merge remote-tracking branch 'origin/feat/nb-egm' into feat/nested-nbegm-ez

Author: hmgaudecker

src/_lcm/egm/continuation.py

@@ -1787,14 +1787,14 @@ def _is_stochastic(
         else:
             row_values = ()
             row_block_shape = ()
-        # A NEGM child carries a stacked candidate axis (keeper + one per outer
-        # node). Detected structurally off the solver spec — `outer_grid` is the
-        # NEGM outer margin's grid — because importing the solver class here
-        # would cycle through the kernel-builder modules that import this one.
-        outer_grid = getattr(target_regime.solver, "outer_grid", None)
-        n_outer_candidates = (
-            int(outer_grid.to_jax().shape[0]) + 1 if outer_grid is not None else 0
-        )
+        # The candidate-axis length comes from the child solver's own
+        # declaration: only a solver whose published carry actually stacks
+        # outer candidates (NEGM's keeper + per-node rows) reports a nonzero
+        # count. Merely having an outer margin is not enough — a solver that
+        # folds its outer axis before publishing (the nested N-NB-EGM upper
+        # envelope) carries no candidate axis, and broadcasting queries over a
+        # phantom axis would desync rows and queries in the read.
+        n_outer_candidates = target_regime.solver.n_stacked_carry_candidates
         reads[target] = _ChildRead(
             next_state_func=get_next_state_function_for_solution(
                 transitions=transitions[target],

src/_lcm/solution/contract.py

@@ -319,6 +319,20 @@ def carry_rows_share_state_grid(self) -> bool:
         """
         return False
 
+    @property
+    def n_stacked_carry_candidates(self) -> int:
+        """Length of the published carry's stacked outer-candidate axis.
+
+        A solver that publishes one carry row per outer durable candidate
+        (keeper plus one per outer-grid node), stacked on an axis before the
+        grid axis, declares that axis length here; a reading parent broadcasts
+        its queries over exactly that many candidates and collapses them by
+        the hard max. A solver whose carry has no candidate axis — no outer
+        margin, or an outer margin already folded inside the solve — declares
+        `0`, and the parent queries each carry row once.
+        """
+        return 0
+
     def validate(self, *, context: SolverBuildContext) -> None:  # noqa: B027
         """Check the regime is in scope for this solver. Default: no-op."""
 

src/_lcm/solution/solvers.py

@@ -484,6 +484,11 @@ def requires_continuation_carries(self) -> bool:
         """NEGM nests a DC-EGM solve that inverts the Euler equation."""
         return True
 
+    @property
+    def n_stacked_carry_candidates(self) -> int:
+        """The published carry stacks the keeper plus one row per outer node."""
+        return int(self.outer_grid.to_jax().shape[0]) + 1
+
     def build_period_kernels(self, *, context: SolverBuildContext) -> SolutionKernels:
         """Build one NEGM period adapter per period, wrapping the inner kernels.
 

tests/solution/test_continuation_stacked_read.py

@@ -18,7 +18,9 @@
     _collapse_stacked_candidates,
 )
 from _lcm.egm.interp import interp_on_padded_grid, prepare_padded_grid
+from lcm.solvers import GridSearch
 from tests.conftest import X64_ENABLED
+from tests.test_models import negm_kinked_toy as negm_toy
 
 # The host comparison reruns the identical interpolation arithmetic, so the two
 # sides agree to the active float precision's roundoff, not better.
@@ -466,3 +468,26 @@ def test_a_poisoned_candidate_at_a_nonzero_index_keeps_the_collapse_nan():
 
     assert bool(jnp.isnan(value[0]))
     assert bool(jnp.isnan(marginal[0]))
+
+
+def test_negm_solver_declares_its_stacked_candidate_count():
+    """A NEGM solver's published carry stacks one row per outer node plus keeper.
+
+    The parent's continuation read sizes its candidate axis off this
+    declaration, so it must equal the outer-grid length plus the keeper slot.
+    """
+    assert (
+        negm_toy.NEGM_SOLVER.n_stacked_carry_candidates
+        == negm_toy.OUTER_GRID.n_points + 1
+    )
+
+
+def test_non_stacking_solvers_declare_zero_stacked_candidates():
+    """A solver that publishes an already-collapsed carry declares no candidates.
+
+    An outer margin folded inside the solver (or no outer margin at all)
+    leaves the carry without a candidate axis; the parent read must then
+    query each carry row once instead of broadcasting queries over a
+    phantom axis.
+    """
+    assert GridSearch().n_stacked_carry_candidates == 0