Make the degenerate-row interpolation contract hold under autodiff

The singleton/empty-row overrides selected the right primal values but ran
after the bracket arithmetic had touched the NaN padding, and `jnp.where`
propagates cotangents through both of its branches — so `jax.grad` of the
read returned NaN (0 · NaN) where the constant clamp's derivative is zero,
poisoning the asset-row Euler marginal whenever a child carry row has one
finite node. The gathered bracket endpoints are now sanitized to finite
dummies on degenerate rows before the arithmetic; the overrides still
publish the contract values. A `lax.cond` dispatch would not fix this: the
production readers vmap the row read, and `cond` lowers to `select` under
`vmap`, re-evaluating both branches.

Also re-pin NaN queries to NaN (the singleton constant would otherwise mask
an upstream failure with a finite value) and add a `Phased` cost-free-base
regression for the composed NEGM resources.

Regressions: query gradient and node-value gradient of a singleton row,
the vmap+grad composition, and the production stacked read differentiated
with a singleton winning candidate.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0123hk3uzBgBeZown4eBBQ2C

Author: hmgaudecker

src/_lcm/egm/interp.py

@@ -151,23 +151,39 @@ def interp_on_prepared_grid(
         jnp.maximum(valid_length - 1, 1),
     ).astype(jnp.int32)
     lower = upper - 1
+    # Degenerate valid prefixes (fewer than two nodes) gather NaN padding into
+    # the bracket. `jnp.where` propagates cotangents through BOTH of its
+    # branches, so NaN partials in the discarded bracket arithmetic would
+    # poison the selected constant's derivative (`0 · NaN = NaN`) — and the
+    # readers are differentiated (asset-row mode grads the continuation read).
+    # Feed the arithmetic a finite dummy bracket on those rows; the overrides
+    # below still publish the contract values.
+    degenerate = valid_length < 2  # noqa: PLR2004
+
+    def _sanitized(gathered: FloatND, dummy: float) -> FloatND:
+        return jnp.where(degenerate, dummy, gathered)
+
     result = _interp_between_nodes(
         x_query=x_query,
-        xp_lower=xp[lower],
-        xp_upper=xp[upper],
-        fp_lower=fp[lower],
-        fp_upper=fp[upper],
-        slope_lower=None if fp_slopes is None else fp_slopes[lower],
-        slope_upper=None if fp_slopes is None else fp_slopes[upper],
+        xp_lower=_sanitized(xp[lower], 0.0),
+        xp_upper=_sanitized(xp[upper], 1.0),
+        fp_lower=_sanitized(fp[lower], 0.0),
+        fp_upper=_sanitized(fp[upper], 0.0),
+        slope_lower=None if fp_slopes is None else _sanitized(fp_slopes[lower], 0.0),
+        slope_upper=None if fp_slopes is None else _sanitized(fp_slopes[upper], 0.0),
     )
-    # Degenerate valid prefixes bypass the bracket arithmetic (whose gathered
-    # upper node is NaN there, which the zero-weight guards would silently turn
-    # into 0.0):
-    # - one valid node ⇒ the edge clamp on both sides is that node's value,
+    # Degenerate-row contract:
+    # - one valid node ⇒ the edge clamp on both sides is that node's value
+    #   (constant in the query, identity in the node's value — also under
+    #   autodiff),
     # - an empty prefix (only from an already-poisoned carry) ⇒ NaN, so the
     #   runtime NaN diagnostics surface the poison instead of a finite constant.
     result = jnp.where(valid_length == 1, fp[0], result)
-    return jnp.where(valid_length == 0, jnp.nan, result)
+    result = jnp.where(valid_length == 0, jnp.nan, result)
+    # A NaN query marks an upstream failure. Regular rows propagate it through
+    # the bracket arithmetic; the degenerate-row constants above would mask it,
+    # so re-pin it explicitly — fail-loud on every row shape.
+    return jnp.where(jnp.isnan(x_query), jnp.nan, result)
 
 
 def interp_right_germ_on_padded_grid(

tests/solution/test_continuation_stacked_read.py

@@ -177,6 +177,50 @@ def test_stacked_read_propagates_a_poisoned_candidate_row():
     assert bool(jnp.isnan(smoothed_value))
 
 
+def test_stacked_read_with_a_singleton_candidate_is_differentiable():
+    """A singleton candidate keeps the read differentiable in the query.
+
+    Asset-row mode publishes the Euler marginal by differentiating the
+    continuation read, so a one-node candidate carry — a constant clamp in the
+    query — must contribute a zero tangent, not a NaN one leaked from its
+    padded bracket. The singleton's clamp value wins here, so the read's query
+    derivative is exactly zero.
+    """
+    singleton = jnp.array([2.0, jnp.nan, jnp.nan])
+    live = jnp.array([0.0, 5.0, 10.0])
+    carry = EGMCarry(
+        endog_grid=jnp.stack([singleton, live])[None, :, :],
+        value=jnp.stack(
+            [jnp.array([10.0, jnp.nan, jnp.nan]), jnp.array([1.0, 2.0, 3.0])]
+        )[None, :, :],
+        marginal_utility=jnp.stack(
+            [jnp.array([0.0, jnp.nan, jnp.nan]), jnp.array([0.5, 0.5, 0.5])]
+        )[None, :, :],
+        taste_shock_scale=jnp.asarray(0.0),
+    )
+    prepared_search_grid, prepared_valid_length = _prepare(carry)
+
+    def read_value(query):
+        smoothed_value, _ = _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=query[None],
+            row_gradients=jnp.asarray([1.0]),
+            n_outer_candidates=2,
+        )
+        return smoothed_value
+
+    value, derivative = jax.value_and_grad(read_value)(jnp.asarray(4.0))
+
+    np.testing.assert_allclose(float(value), 10.0, atol=_READ_ATOL)
+    np.testing.assert_allclose(float(derivative), 0.0, atol=_READ_ATOL)
+
+
 def test_stacked_read_tie_publishes_the_right_continuous_marginal():
     """At an exact candidate crossing the production read selects the right winner.
 

tests/solution/test_egm_interp.py

@@ -145,6 +145,93 @@ def test_singleton_row_clamps_to_its_node():
     np.testing.assert_array_equal(np.asarray(marginal), [2.0, 2.0, 2.0])
 
 
+def test_singleton_row_is_constant_under_autodiff():
+    """A singleton row's constant clamp differentiates as a constant.
+
+    `jax.grad` with respect to the query must be exactly zero (finite, not
+    NaN): asset-row mode differentiates the continuation read in the Euler
+    slot, so a NaN tangent leaking from the padded bracket would poison the
+    published Euler marginal even though the primal value is correct.
+    """
+    xp = jnp.array([1.0, jnp.nan, jnp.nan])
+    fp = jnp.array([5.0, jnp.nan, jnp.nan])
+    slopes = jnp.array([2.0, jnp.nan, jnp.nan])
+
+    def read(query):
+        return interp.interp_on_padded_grid(
+            x_query=query, xp=xp, fp=fp, fp_slopes=slopes
+        )
+
+    value, derivative = jax.value_and_grad(read)(jnp.asarray(1.0))
+
+    np.testing.assert_allclose(float(value), 5.0, atol=1e-12)
+    np.testing.assert_allclose(float(derivative), 0.0, atol=1e-12)
+
+
+def test_singleton_row_passes_its_node_values_tangent_through():
+    """The singleton clamp is the node's value: unit derivative in that node.
+
+    Differentiating the read with respect to the row's single valid value must
+    give exactly one — the clamp publishes that value verbatim — with no NaN
+    contamination from the padded slots.
+    """
+    xp = jnp.array([1.0, jnp.nan, jnp.nan])
+    slopes = jnp.array([2.0, jnp.nan, jnp.nan])
+
+    def read(node_value):
+        fp = jnp.array([jnp.nan, jnp.nan, jnp.nan]).at[0].set(node_value)
+        return interp.interp_on_padded_grid(
+            x_query=jnp.asarray(2.0), xp=xp, fp=fp, fp_slopes=slopes
+        )
+
+    derivative = jax.grad(read)(jnp.asarray(5.0))
+
+    np.testing.assert_allclose(float(derivative), 1.0, atol=1e-12)
+
+
+def test_singleton_row_stays_constant_under_vmap_and_autodiff():
+    """The zero query derivative survives per-row `vmap` batching.
+
+    The production readers map the row read over a stacked candidate axis, so
+    the degenerate-row handling must stay AD-safe inside `vmap` — a construct
+    that re-evaluates both sides of a branch (as `lax.cond` lowered under
+    `vmap` would) leaks the padded bracket's NaN tangent back in.
+    """
+    xp = jnp.array([[1.0, 2.0, 3.0], [1.0, jnp.nan, jnp.nan]])
+    fp = jnp.array([[1.0, 4.0, 9.0], [5.0, jnp.nan, jnp.nan]])
+    slopes = jnp.array([[2.0, 4.0, 6.0], [2.0, jnp.nan, jnp.nan]])
+
+    def summed_read(query):
+        def read_row(xp_row, fp_row, slopes_row):
+            return interp.interp_on_padded_grid(
+                x_query=query, xp=xp_row, fp=fp_row, fp_slopes=slopes_row
+            )
+
+        return jnp.sum(jax.vmap(read_row)(xp, fp, slopes))
+
+    derivative = jax.grad(summed_read)(jnp.asarray(2.5))
+
+    # The regular row holds `x²` data, which the Hermite read reproduces
+    # (derivative `2q = 5`); the singleton row contributes exactly zero.
+    np.testing.assert_allclose(float(derivative), 5.0, atol=1e-12)
+
+
+def test_nan_query_reads_nan_on_a_singleton_row():
+    """A NaN query stays NaN even where the singleton clamp is constant.
+
+    A NaN query marks an upstream failure; the constant clamp must not convert
+    it into a finite value — regular rows already propagate NaN queries through
+    the bracket arithmetic, and singleton rows must fail just as loudly.
+    """
+    xp = jnp.array([1.0, jnp.nan, jnp.nan])
+    fp = jnp.array([5.0, jnp.nan, jnp.nan])
+
+    got = interp.interp_on_padded_grid(x_query=jnp.array([jnp.nan, 1.0]), xp=xp, fp=fp)
+
+    assert bool(jnp.isnan(got[0]))
+    np.testing.assert_allclose(float(got[1]), 5.0, atol=1e-12)
+
+
 def test_empty_row_reads_nan():
     """An all-NaN (poisoned) row reads NaN at every query, never a finite value.
 

tests/solution/test_negm_validation.py

@@ -27,7 +27,7 @@
     validate_negm_regimes,
 )
 from _lcm.regime_building.finalize import finalize_regimes
-from lcm import DiscreteGrid, ExtremeValueTasteShocks, categorical
+from lcm import DiscreteGrid, ExtremeValueTasteShocks, Phased, categorical
 from lcm.exceptions import ModelInitializationError
 from lcm.regime import Regime as UserRegime
 from lcm.typing import (
@@ -366,6 +366,35 @@ def test_finalize_composes_resources_as_base_minus_outer_cost():
     ) == pytest.approx(5.0)
 
 
+def test_finalize_composes_resources_with_a_phased_base():
+    """A `Phased` cost-free base still yields the composed resources function.
+
+    The synthesized resources function reads the base and the cost by name, so
+    phase resolution happens at the producer level: composition succeeds with a
+    `Phased` base slot and the injected function is the plain difference of its
+    two inputs in either phase.
+    """
+    regime = _VALID.replace(
+        functions={
+            **dict(_VALID.functions),
+            "resources_before_outer_cost": Phased(
+                solve=_cost_free_base, simulate=_cost_free_base
+            ),
+        },
+    )
+
+    finalized = finalize_regimes(
+        user_regimes={"alive": regime}, derived_categoricals={}
+    )["alive"]
+    composed = cast("UserFunction", finalized.functions["resources"])
+
+    assert float(
+        composed(
+            resources_before_outer_cost=jnp.asarray(7.0), credited=jnp.asarray(2.0)
+        )
+    ) == pytest.approx(5.0)
+
+
 def test_missing_outer_cost_with_costful_resources_is_rejected():
     """Omitting `NEGM.outer_cost` while resources reads the outer margin fails.