Skip to content

Commit f83fac0

Browse files
hmgaudeckerclaude
andcommitted
Merge feat/dcegm: left-germ terminal tie ownership + below-support secant JVP
Brings in the round-14 F3/F4 repairs (da5fd88): at an exact shared-upper terminal tie the stacked collapse decides ownership by the complete left germ (the branch that wins on the supported left neighborhood owns the published marginal), and the custom JVP's query tangent below support is the first finite bracket's secant — the derivative of the primal read's documented extension — instead of a hard zero. Resolution notes: union imports; doc-variant hunks keep this branch's wording; the poison-collapse regression passes the new left-germ argument; the germ computation and support mask move to a module-level helper (`_candidate_germs_and_support_mask`) with a typed germ reshape. Both round-14 counterexample tests (terminal-tie left owner, below-support custom JVP) pass verbatim on the merged tree. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01BNQkLG5QzXcgGdpgH6sN4u
2 parents e57ca73 + da5fd88 commit f83fac0

8 files changed

Lines changed: 565 additions & 97 deletions

src/_lcm/egm/continuation.py

Lines changed: 120 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
)
3333
from _lcm.egm.interp import (
3434
interp_and_derivative_on_prepared_grid,
35+
interp_left_germ_on_prepared_grid,
3536
interp_on_prepared_grid,
3637
interp_right_germ_on_prepared_grid,
3738
locate_on_grid,
@@ -1437,44 +1438,17 @@ def value_and_slope_row(
14371438
search_rows, valid_rows, grid_rows, marginal_rows, queries_flat
14381439
)
14391440
if n_outer_candidates:
1440-
# Below a candidate's own first finite coh node its support has not
1441-
# started: mask the read to `-inf` so the edge clamp cannot hand an
1442-
# infeasible lifted candidate a boundary value that wins the max. The
1443-
# `-inf` also pins the marginal to zero below. The right germ of each
1444-
# candidate's value read feeds the tie rule at the candidate max; it
1445-
# needs no support mask of its own — a below-support candidate enters
1446-
# the tie set only when every candidate is below support, where all
1447-
# published marginals are exactly zero regardless of the winner.
1448-
def right_germ_row(
1449-
search_grid: Float1D,
1450-
valid_length: ScalarInt,
1451-
xp: Float1D,
1452-
fp: Float1D,
1453-
fp_slopes: Float1D,
1454-
x_query: ScalarFloat,
1455-
) -> tuple[ScalarBool, ScalarFloat, ScalarFloat, ScalarFloat]:
1456-
"""Right germ of one value row at its query; positional per `jax.vmap`."""
1457-
return interp_right_germ_on_prepared_grid(
1458-
x_query=x_query,
1459-
search_grid=search_grid,
1460-
valid_length=valid_length,
1461-
xp=xp,
1462-
fp=fp,
1463-
fp_slopes=fp_slopes,
1441+
right_germ_at_child, left_germ_at_child, value_at_child = (
1442+
_candidate_germs_and_support_mask(
1443+
search_rows=search_rows,
1444+
valid_rows=valid_rows,
1445+
grid_rows=grid_rows,
1446+
value_rows=value_rows,
1447+
marginal_rows=marginal_rows,
1448+
queries_flat=queries_flat,
1449+
value_at_child=value_at_child,
14641450
)
1465-
1466-
right_germ_at_child = jax.vmap(right_germ_row)(
1467-
search_rows, valid_rows, grid_rows, value_rows, marginal_rows, queries_flat
1468-
)
1469-
row_lower = jnp.min(
1470-
jnp.where(jnp.isfinite(grid_rows), grid_rows, jnp.inf), axis=1
14711451
)
1472-
# Only rows with a finite first node have a support to fall below;
1473-
# `row_lower` is `+inf` on an all-NaN (poisoned) row, whose NaN read
1474-
# must reach the candidate maximum fail-loud instead of becoming an
1475-
# ordinary infeasible `(-inf, 0)` pair.
1476-
below_row_support = (queries_flat < row_lower) & jnp.isfinite(row_lower)
1477-
value_at_child = jnp.where(below_row_support, -jnp.inf, value_at_child)
14781452
# `-inf` entries interpolate pointwise to `-inf` (never NaN) and carry
14791453
# exactly-zero marginal utility, so an infeasible-everywhere row reads as
14801454
# the `-inf` / zero pair while a row with isolated `-inf` nodes (e.g. a
@@ -1494,8 +1468,11 @@ def right_germ_row(
14941468
value_at_child, marginal_at_child = _collapse_stacked_candidates(
14951469
value_at_child=value_at_child,
14961470
marginal_at_child=marginal_at_child,
1497-
right_germ_at_child=tuple(
1498-
component.reshape(block_shape) for component in right_germ_at_child
1471+
right_germ_at_child=_reshape_germ(
1472+
germ=right_germ_at_child, block_shape=block_shape
1473+
),
1474+
left_germ_at_child=_reshape_germ(
1475+
germ=left_germ_at_child, block_shape=block_shape
14991476
),
15001477
)
15011478

@@ -1572,6 +1549,96 @@ def _blend_passive_axes(
15721549
return value_at_child, marginal_at_child
15731550

15741551

1552+
def _candidate_germs_and_support_mask(
1553+
*,
1554+
search_rows: FloatND,
1555+
valid_rows: IntND,
1556+
grid_rows: FloatND,
1557+
value_rows: FloatND,
1558+
marginal_rows: FloatND,
1559+
queries_flat: Float1D,
1560+
value_at_child: Float1D,
1561+
) -> tuple[
1562+
tuple[BoolND, FloatND, FloatND, FloatND],
1563+
tuple[BoolND, FloatND, FloatND, FloatND],
1564+
Float1D,
1565+
]:
1566+
"""Compute per-candidate value-read germs and mask reads below support.
1567+
1568+
Below a candidate's own first finite coh node its support has not
1569+
started: the read is masked to `-inf` so the edge clamp cannot hand an
1570+
infeasible lifted candidate a boundary value that wins the max (the
1571+
`-inf` also pins the marginal to zero below). The germs of each
1572+
candidate's value read feed the tie rule at the candidate max; they need
1573+
no support mask of their own — the left germ is dead at or below the
1574+
first finite node by construction, and a below-support candidate enters
1575+
the tie set only when every candidate is below support, where all
1576+
published marginals are exactly zero regardless of the winner.
1577+
1578+
Returns:
1579+
Tuple of the right germs, the left germs, and the support-masked
1580+
value reads, all aligned with the flat candidate-row axis.
1581+
1582+
"""
1583+
1584+
def germ_rows(
1585+
search_grid: Float1D,
1586+
valid_length: ScalarInt,
1587+
xp: Float1D,
1588+
fp: Float1D,
1589+
fp_slopes: Float1D,
1590+
x_query: ScalarFloat,
1591+
) -> tuple[
1592+
tuple[ScalarBool, ScalarFloat, ScalarFloat, ScalarFloat],
1593+
tuple[ScalarBool, ScalarFloat, ScalarFloat, ScalarFloat],
1594+
]:
1595+
"""Right and left germs of one value row at its query; per `jax.vmap`."""
1596+
right_germ = interp_right_germ_on_prepared_grid(
1597+
x_query=x_query,
1598+
search_grid=search_grid,
1599+
valid_length=valid_length,
1600+
xp=xp,
1601+
fp=fp,
1602+
fp_slopes=fp_slopes,
1603+
)
1604+
left_germ = interp_left_germ_on_prepared_grid(
1605+
x_query=x_query,
1606+
search_grid=search_grid,
1607+
valid_length=valid_length,
1608+
xp=xp,
1609+
fp=fp,
1610+
fp_slopes=fp_slopes,
1611+
)
1612+
return right_germ, left_germ
1613+
1614+
right_germ_at_child, left_germ_at_child = jax.vmap(germ_rows)(
1615+
search_rows, valid_rows, grid_rows, value_rows, marginal_rows, queries_flat
1616+
)
1617+
row_lower = jnp.min(jnp.where(jnp.isfinite(grid_rows), grid_rows, jnp.inf), axis=1)
1618+
# Only rows with a finite first node have a support to fall below;
1619+
# `row_lower` is `+inf` on an all-NaN (poisoned) row, whose NaN read
1620+
# must reach the candidate maximum fail-loud instead of becoming an
1621+
# ordinary infeasible `(-inf, 0)` pair.
1622+
below_row_support = (queries_flat < row_lower) & jnp.isfinite(row_lower)
1623+
value_at_child = jnp.where(below_row_support, -jnp.inf, value_at_child)
1624+
return right_germ_at_child, left_germ_at_child, value_at_child
1625+
1626+
1627+
def _reshape_germ(
1628+
*,
1629+
germ: tuple[BoolND, FloatND, FloatND, FloatND],
1630+
block_shape: tuple[int, ...],
1631+
) -> tuple[BoolND, FloatND, FloatND, FloatND]:
1632+
"""Reshape each germ component from the flat row axis to the row block."""
1633+
finite_flag, first, second, third = germ
1634+
return (
1635+
finite_flag.reshape(block_shape),
1636+
first.reshape(block_shape),
1637+
second.reshape(block_shape),
1638+
third.reshape(block_shape),
1639+
)
1640+
1641+
15751642
def _fail_if_carry_shape_mismatches_declaration(
15761643
*,
15771644
block_shape: tuple[int, ...],
@@ -1606,6 +1673,7 @@ def _collapse_stacked_candidates(
16061673
value_at_child: FloatND,
16071674
marginal_at_child: FloatND,
16081675
right_germ_at_child: tuple[BoolND, FloatND, FloatND, FloatND],
1676+
left_germ_at_child: tuple[BoolND, FloatND, FloatND, FloatND],
16091677
) -> tuple[FloatND, FloatND]:
16101678
"""Collapse the candidate axis by the exact hard max at the query.
16111679
@@ -1614,10 +1682,13 @@ def _collapse_stacked_candidates(
16141682
tie resolves right-continuously in the value read itself: the tied
16151683
candidates are compared by the complete right germ of their own value
16161684
interpolants (`right_germ_winner` — right-finiteness, then the first three
1617-
one-sided derivatives, then the lowest index among locally identical
1618-
pieces), so the branch whose read actually wins immediately to the right
1619-
owns the (economic) marginal the parent's Euler inversion consumes. (The
1620-
germ uses the unscaled interpolant derivatives: the composed gradient is
1685+
one-sided derivatives), so the branch whose read actually wins immediately
1686+
to the right owns the (economic) marginal the parent's Euler inversion
1687+
consumes; right-identical candidates (a shared terminal abscissa, where
1688+
every candidate clamps) are separated by their left germs so the marginal
1689+
stays inside the envelope's generalized gradient at the boundary, and only
1690+
branches identical on both sides fall back to the lowest index. (The
1691+
germs use the unscaled interpolant derivatives: the composed gradient is
16211692
shared by all candidates of a cell and positive, so scaling could never
16221693
reorder them.) A cell whose candidates are all `-inf` (no live support)
16231694
keeps the `(-inf, 0)` infeasible contract: every masked marginal is
@@ -1628,13 +1699,19 @@ def _collapse_stacked_candidates(
16281699
marginal_at_child: Gradient-scaled candidate marginal reads, same shape.
16291700
right_germ_at_child: Tuple of the right-finiteness flag and the first
16301701
three right derivatives of the candidate value reads, same shape.
1702+
left_germ_at_child: Tuple of the left-finiteness flag and the first
1703+
three left derivatives of the candidate value reads, same shape.
16311704
16321705
Returns:
16331706
Tuple of the winner's value and marginal, with the candidate axis
16341707
collapsed.
16351708
16361709
"""
1637-
winner = right_germ_winner(value=value_at_child, right_germ=right_germ_at_child)
1710+
winner = right_germ_winner(
1711+
value=value_at_child,
1712+
right_germ=right_germ_at_child,
1713+
left_germ=left_germ_at_child,
1714+
)
16381715
winner_marginal = jnp.take_along_axis(marginal_at_child, winner, axis=-1)[..., 0]
16391716
# The published value is the maximum itself: identical to the winner's read
16401717
# at any tie, and NaN-propagating when a poisoned candidate row (whose NaN

0 commit comments

Comments
 (0)