Exact MSS ties + same-point crossing residual + read-contract hardenings

Round-14 external-audit repairs, all confirmed by execution first:

- `_lexicographic_winner` publishes the exact stored maximum as the side
  value (the tie winner owns policy and slope convention only), and value
  ties are exact stored equality — the only translation-invariant tie set.
  `same_record` likewise compares exactly (one genuine record is bitwise
  equal). A tolerance window on absolute magnitudes let a common cardinal
  shift hand the node to a near-maximal competitor's lower record.
- A crossing tie needs more than offset proximity: a line joins the group
  only if it passes through the same numerical point (value residual at
  the emitted crossing within the cancellation-error scale of the locally
  computed magnitudes). Distinct crossings at large local offsets are
  enumerated sequentially instead of being handed to the steepest line.
- The distributed-fold spec compares with a mixed relative/scale-anchored
  bound (values can sit near zero, where a pure-relative few-ulp bound is
  spuriously strict) after asserting identical finite masks.
- The child read checks the carry block's leading shape against the
  solver's stacked-candidate declaration before broadcasting and names the
  violated contract, instead of failing as a vmap axis-size mismatch.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01BNQkLG5QzXcgGdpgH6sN4u

Author: hmgaudecker

src/_lcm/egm/continuation.py

@@ -1331,6 +1331,11 @@ def _aggregate_child_choices(
     # Leading axes of the blocks: the child's passive nodes, then its
     # discrete-action combos (then the candidate axis of a stacked NEGM child).
     block_shape = value_block.shape[:-1]
+    _fail_if_carry_shape_mismatches_declaration(
+        block_shape=block_shape,
+        row_queries_shape=row_queries.shape,
+        n_outer_candidates=n_outer_candidates,
+    )
     if n_outer_candidates:
         # The stacked candidates share the lifted common-coh axis, so every
         # candidate row of a block cell is read at that cell's single query and
@@ -1567,6 +1572,35 @@ def _blend_passive_axes(
     return value_at_child, marginal_at_child
 
 
+def _fail_if_carry_shape_mismatches_declaration(
+    *,
+    block_shape: tuple[int, ...],
+    row_queries_shape: tuple[int, ...],
+    n_outer_candidates: int,
+) -> None:
+    """Check the carry block's leading shape against the stacking declaration.
+
+    The declaration and the published carry must agree before any
+    broadcasting: a mismatch would otherwise surface as an opaque vmap
+    axis-size error deep in the batched interpolation instead of naming the
+    violated solver contract.
+    """
+    expected_block_shape = (
+        (*row_queries_shape, n_outer_candidates)
+        if n_outer_candidates
+        else row_queries_shape
+    )
+    if block_shape != expected_block_shape:
+        msg = (
+            "The child's published carry block has leading shape "
+            f"{block_shape}, but its solver declares "
+            f"n_stacked_carry_candidates={n_outer_candidates}, which requires "
+            f"{expected_block_shape}. The solver's declaration must match the "
+            "candidate-axis structure of the carry it publishes."
+        )
+        raise ValueError(msg)
+
+
 def _collapse_stacked_candidates(
     *,
     value_at_child: FloatND,

src/_lcm/egm/upper_envelope/mss.py

@@ -214,18 +214,18 @@ def refine_envelope_with_support(
         [jnp.full((1,), -jnp.inf, dtype=query_grid.dtype), query_grid[:-1]]
     )
     is_new = query_grid > prev_grid
-    # Record equality is a few-ulp window: two computations of the same
-    # quantity agree to within a handful of rounding steps of the stored
-    # magnitude, and gaps below that scale are indistinguishable in storage.
-    # A wider cushion would swallow genuinely representable gaps whenever a
-    # common cardinal shift raises the value level — the winner would then
-    # depend on an arbitrary value normalization. A missed tie merely
-    # duplicates the node, which the read resolves by its side convention.
+    # Record equality is exact: when both side winners are the same branch,
+    # its value and policy at the node come from the same link evaluated at
+    # the same query — identical arithmetic, so genuinely one record is
+    # bitwise equal. Any tolerance window here scales with the absolute value
+    # level and would merge representable records under a common cardinal
+    # shift. A missed merge merely duplicates the node, which the read
+    # resolves by its side convention.
     tolerance = 8.0 * float(jnp.finfo(query_grid.dtype).eps)
     same_record = (
         (left_side.branch == right_side.branch)
-        & jnp.isclose(left_side.value, right_side.value, rtol=tolerance, atol=0.0)
-        & jnp.isclose(left_side.policy, right_side.policy, rtol=tolerance, atol=0.0)
+        & (left_side.value == right_side.value)
+        & (left_side.policy == right_side.policy)
     )
     node_left_valid = left_side.exists & is_new
     node_right_valid = right_side.exists & is_new & ~(left_side.exists & same_record)
@@ -493,23 +493,25 @@ def _lexicographic_winner(
     masked_value = jnp.where(covers, value_at, -jnp.inf)
     best_value = jnp.max(masked_value, axis=1)
     exists = jnp.isfinite(best_value)
-    # Value ties are a few-ulp window of the stored magnitude — the scale at
-    # which two computations of the same envelope value are indistinguishable.
-    # A wider cushion would let a common cardinal value shift merge branches
-    # whose gap is genuinely representable. A missed tie merely skips the
+    # A value tie is exact stored equality — the only translation-invariant
+    # tie set: any tolerance window scales with the absolute value level, so
+    # a common cardinal shift would pull genuinely representable gaps inside
+    # it and let a lower record own the node. A missed tie (two computations
+    # of the same crossing value a rounding step apart) merely skips the
     # slope tie-break and duplicates the node downstream — never merges
     # records.
-    tolerance = 8.0 * jnp.finfo(value_at.dtype).eps
-    tie = covers & jnp.isclose(
-        masked_value, best_value[:, None], rtol=tolerance, atol=0.0
-    )
+    tie = covers & (masked_value == best_value[:, None])
     slope_score = jnp.where(tie, slope_sign * links.value_slope[None, :], -jnp.inf)
     link = jnp.argmax(slope_score, axis=1).astype(jnp.int32)
+    # The published value is the exact maximum itself, not the tie winner's
+    # own read: the winner owns the policy and the slope convention, but a
+    # numerical tie rule must never replace the hard maximum of the stored
+    # records by a near-maximal competitor's lower value.
     return _SideWinner(
         exists=exists,
         link=link,
         branch=jnp.where(exists, links.segment[link], -1).astype(jnp.int32),
-        value=jnp.take_along_axis(value_at, link[:, None], axis=1)[:, 0],
+        value=best_value,
         policy=jnp.take_along_axis(policy_at, link[:, None], axis=1)[:, 0],
     )
 
@@ -617,7 +619,28 @@ def step(
         offset_next = jnp.min(offset_masked, axis=1)
         found = jnp.isfinite(offset_next)
         offset_tolerance = crossing_ulp * jnp.where(found, offset_next, 0.0)
-        tie = valid & (offset <= (offset_next + offset_tolerance)[:, None])
+        near_minimal = valid & (offset <= (offset_next + offset_tolerance)[:, None])
+        # Offset proximity alone cannot certify a simultaneous crossing: the
+        # window scales with the offset itself, so inside it two crossings can
+        # sit many representable positions apart, and handing the group to the
+        # steepest line would skip the branch that wins between them. A line
+        # joins the tie only if it also passes through the same numerical
+        # point — its value at the crossing agrees with the winner line's to a
+        # few ulp of the locally computed magnitudes (the cancellation-error
+        # scale of evaluating the difference). A genuinely distinct crossing
+        # fails the residual test and is enumerated on a later step instead.
+        x_next_safe = x_current + jnp.where(found, offset_next, 0.0)
+        value_at_next = links.value_at(x_next_safe[:, None])
+        winner_value_next = jnp.take_along_axis(value_at_next, winner[:, None], axis=1)
+        residual_scale = jnp.maximum(jnp.abs(value_at_next), jnp.abs(winner_value_next))
+        same_point = jnp.abs(value_at_next - winner_value_next) <= (
+            8.0 * jnp.finfo(query_grid.dtype).eps * residual_scale
+        )
+        # The exact-minimum line(s) define the crossing being emitted, so they
+        # belong to the tie unconditionally — the residual test only decides
+        # which *other* near-minimal lines genuinely pass through the point.
+        is_earliest = valid & (offset <= offset_next[:, None])
+        tie = (near_minimal & same_point) | is_earliest
         incoming = jnp.argmax(
             jnp.where(tie, links.value_slope[None, :], -jnp.inf), axis=1
         ).astype(jnp.int32)

tests/solution/test_continuation_stacked_read.py

@@ -11,6 +11,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
+import pytest
 
 from _lcm.egm.carry import EGMCarry
 from _lcm.egm.continuation import (
@@ -491,3 +492,34 @@ def test_non_stacking_solvers_declare_zero_stacked_candidates():
     phantom axis.
     """
     assert GridSearch().n_stacked_carry_candidates == 0
+
+
+def test_a_carry_without_the_declared_candidate_axis_fails_with_a_contract_error():
+    """A stacking declaration that disagrees with the carry shape names itself.
+
+    When the declared candidate count is positive but the carry block carries
+    no matching candidate axis, the read must raise a solver-contract error
+    that names the declaration — not fall through to an opaque vmap axis-size
+    mismatch deep in the batched interpolation.
+    """
+    carry = EGMCarry(
+        endog_grid=jnp.broadcast_to(jnp.array([0.0, 5.0, 10.0]), (2, 3)),
+        value=jnp.ones((2, 3)),
+        marginal_utility=jnp.zeros((2, 3)),
+        taste_shock_scale=jnp.asarray(0.0),
+    )
+    prepared_search_grid, prepared_valid_length = _prepare(carry)
+
+    with pytest.raises(ValueError, match="n_stacked_carry_candidates"):
+        _aggregate_child_choices(
+            carry=carry,
+            prepared_search_grid=prepared_search_grid,
+            prepared_valid_length=prepared_valid_length,
+            has_taste_shocks=False,
+            child_index=(),
+            child_passive_values=(jnp.asarray(0.0),),
+            child_passive_grids=(jnp.asarray([0.0]),),
+            row_queries=jnp.asarray([4.0, 4.0]),
+            row_gradients=jnp.asarray([1.0, 1.0]),
+            n_outer_candidates=2,
+        )

tests/solution/test_mss_invariance_and_support.py

@@ -361,3 +361,113 @@ def test_two_distinct_representable_crossings_are_both_emitted():
     )
     truth = intercepts + slopes * 1.25
     np.testing.assert_allclose(read_policy, float(policies[truth.argmax()]), rtol=1e-5)
+
+
+def _read_two_branch_row_at_origin(shift: float) -> tuple[float, float]:
+    """Refine a two-branch f32 correspondence and read value/policy at the origin.
+
+    Branch values at the origin node differ by four units — a genuinely
+    representable gap at every level used by the callers — so the published
+    read must be the higher branch (value `4 + shift`, policy 1000) no matter
+    what common cardinal level `shift` puts the values at.
+    """
+    origin = np.float32(2000.0)
+    width = np.float32(100.0)
+    policies = np.array([1000.0, 800.0], dtype=np.float32)
+    slopes = np.float32(10000.0) / policies
+    intercepts = np.array([4.0, 0.0], dtype=np.float32) + np.float32(shift)
+
+    grid, policy, value = [], [], []
+    for slope, intercept, action in zip(slopes, intercepts, policies, strict=True):
+        grid.extend([origin, origin + width])
+        policy.extend([action, action])
+        value.extend([intercept, intercept + slope * width])
+
+    refined_grid, refined_policy, refined_value, _ = refine_envelope(
+        endog_grid=jnp.asarray(np.asarray(grid, dtype=np.float32)),
+        policy=jnp.asarray(np.asarray(policy, dtype=np.float32)),
+        value=jnp.asarray(np.asarray(value, dtype=np.float32)),
+        n_refined=16,
+    )
+    query = jnp.asarray(origin, dtype=jnp.float32)
+    got_value = float(
+        interp_on_padded_grid(x_query=query, xp=refined_grid, fp=refined_value)
+    )
+    got_policy = float(
+        interp_on_padded_grid(x_query=query, xp=refined_grid, fp=refined_policy)
+    )
+    return got_value, got_policy
+
+
+def test_a_large_common_value_shift_cannot_lower_the_published_maximum():
+    """The published value is the exact stored maximum at any cardinal level.
+
+    A representable four-unit branch gap survives a common shift of `1e7` in
+    f32 (both shifted values are distinct floats), so the read must publish
+    the higher branch's value and policy — never a near-maximal competitor's
+    lower record.
+    """
+    base_value, base_policy = _read_two_branch_row_at_origin(0.0)
+    shifted_value, shifted_policy = _read_two_branch_row_at_origin(10_000_000.0)
+
+    np.testing.assert_allclose(base_value, 4.0, rtol=0.0, atol=0.0)
+    np.testing.assert_allclose(base_policy, 1000.0, rtol=0.0, atol=0.0)
+    np.testing.assert_allclose(shifted_value, 10_000_004.0, rtol=0.0, atol=0.0)
+    np.testing.assert_allclose(shifted_policy, 1000.0, rtol=0.0, atol=0.0)
+
+
+def test_two_distinct_crossings_at_large_local_offsets_are_both_emitted():
+    """Crossings one unit apart at local offset `1e6` are enumerated in order.
+
+    At that offset f32 spacing is `0.0625`, so the two crossing abscissae are
+    sixteen ulp apart — genuinely distinct points whose middle branch owns the
+    envelope between them. A tie may coalesce crossings only when the lines
+    actually meet at the same numerical point; here the third line sits about
+    twenty ulp of the value scale below the winner at the first crossing, so
+    the scan must emit both switches and the strictly-between read must
+    follow the middle branch.
+    """
+    origin = np.float32(5_000_000.0)
+    width = np.float32(1_000_010.0)
+    policies = np.array([4_000_000.0, 2_500_000.0, 1_000_000.0], dtype=np.float32)
+    slopes = 1.0 / policies
+    first_crossing = np.float32(1_000_000.0)
+    second_crossing = np.float32(1_000_001.0)
+    intercepts = np.array(
+        [
+            0.0,
+            (slopes[0] - slopes[1]) * first_crossing,
+            (slopes[0] - slopes[1]) * first_crossing
+            + (slopes[1] - slopes[2]) * second_crossing,
+        ],
+        dtype=np.float32,
+    )
+
+    grid, policy, value = [], [], []
+    for slope, intercept, action in zip(slopes, intercepts, policies, strict=True):
+        grid.extend([origin, origin + width])
+        policy.extend([action, action])
+        value.extend([intercept, intercept + slope * width])
+
+    refined_grid, refined_policy, refined_value, n_kept = refine_envelope(
+        endog_grid=jnp.asarray(np.asarray(grid, dtype=np.float32)),
+        policy=jnp.asarray(np.asarray(policy, dtype=np.float32)),
+        value=jnp.asarray(np.asarray(value, dtype=np.float32)),
+        n_refined=32,
+    )
+    local_grid = np.asarray(refined_grid)[: int(n_kept)] - origin
+
+    assert np.isclose(local_grid, first_crossing, rtol=0.0, atol=0.25).sum() == 2
+    assert np.isclose(local_grid, second_crossing, rtol=0.0, atol=0.25).sum() == 2
+
+    local_query = np.float32(1_000_000.5)
+    query = jnp.asarray(origin + local_query, dtype=jnp.float32)
+    got_value = float(
+        interp_on_padded_grid(x_query=query, xp=refined_grid, fp=refined_value)
+    )
+    got_policy = float(
+        interp_on_padded_grid(x_query=query, xp=refined_grid, fp=refined_policy)
+    )
+    truth = intercepts + slopes * local_query
+    np.testing.assert_allclose(got_value, float(truth.max()), rtol=1e-5, atol=1e-9)
+    np.testing.assert_allclose(got_policy, float(policies[truth.argmax()]))

tests/solution/test_nbegm_stochastic_distributed_fold.py

@@ -7,7 +7,10 @@
 (2) differ in size, a mis-targeted fold axis is a hard broadcasting error, so the
 distributed solve must both build and reproduce the non-distributed value at
 reassociation level: XLA fuses the read arithmetic differently per axis layout,
-so the two solves agree to a few ulp of the value scale, not bit-for-bit.
+so the two solves agree to a few ulp of the value scale, not bit-for-bit. The
+comparison mixes a relative and a scale-anchored absolute component because
+value entries can sit near zero, where a pure-relative few-ulp bound would be
+spuriously strict.
 """
 
 import numpy as np
@@ -37,4 +40,12 @@ def test_stochastic_fold_is_invariant_to_distributing_a_ride_state():
         # is `(kind, income, liquid)` where the plain solve is `(income, kind,
         # liquid)`; align the two before the exact-value comparison.
         distributed_v = np.moveaxis(np.asarray(distributed[period]["alive"]), 0, 1)
-        np.testing.assert_allclose(distributed_v, plain_v, rtol=1e-15, atol=0.0)
+        np.testing.assert_array_equal(np.isfinite(distributed_v), np.isfinite(plain_v))
+        eps = 8.0 * np.finfo(plain_v.dtype).eps
+        finite = np.isfinite(plain_v)
+        scale = float(
+            max(np.abs(plain_v[finite]).max(), np.abs(distributed_v[finite]).max())
+        )
+        np.testing.assert_allclose(
+            distributed_v[finite], plain_v[finite], rtol=eps, atol=eps * scale
+        )