log_add: shift-guard NaN gradient at structural-zero (-inf) - (-inf)

When both operands are structural zeros (sign=0, log_mag=-inf),
``small_log - big_log`` evaluates to ``(-inf) - (-inf) = NaN``.
The downstream ``where`` cascade masks the forward value (returns
LOG_ZERO correctly), but reverse-mode autodiff still differentiates
through the NaN expression and propagates a NaN gradient.

Apply the same ``big_log_safe = where(isneginf(big_log), 0, big_log)``
shift used in log_sum so the dangerous branches stay finite.  The
both-zero case is already masked to LOG_ZERO regardless, so the
shift is irrelevant to the forward value but unblocks the gradient.

Add test_log_add_both_zero_grad_finite as a regression lock.

Latent fragility — production jet rules go through log_sum (already
safe) — so no acopula experiment results were affected.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: thisiscam

jet_array/log_space_ops.py

@@ -110,14 +110,23 @@ def log_add(a: LogSeries, b: LogSeries) -> LogSeries:
     small_log = jnp.where(a_bigger, b.log_mag, a.log_mag)
     big_sign = jnp.where(a_bigger, a.sign, b.sign)
 
-    # `delta = small_log - big_log <= 0`. Note: -inf - finite = -inf gives
-    # exp(-inf) = 0, so the small=zero case naturally yields log|sum| = big.
-    delta = small_log - big_log
+    # Shift big_log to a finite reference when both operands are structural
+    # zero (big_log = -inf), so that `small_log - big_log_safe = -inf - 0 =
+    # -inf` keeps the dangerous branches finite. The both-zero case is
+    # masked out below to LOG_ZERO regardless, so this shift is irrelevant
+    # to the forward value but keeps the reverse-mode gradient finite.
+    big_log_safe = jnp.where(jnp.isneginf(big_log),
+                             jnp.zeros_like(big_log), big_log)
+
+    # `delta = small_log - big_log_safe <= 0`. With the shift, both-zero
+    # gives delta = -inf instead of NaN; finite-vs-zero unchanged because
+    # big_log was already finite.
+    delta = small_log - big_log_safe
 
     # log1p(exp(delta))   for same sign  (always finite, >= log 1 = 0)
     # log1p(-exp(delta))  for opposite sign (could be -inf at exact cancellation)
-    same_sign_log = big_log + jnp.log1p(jnp.exp(delta))
-    opp_sign_log = big_log + jnp.log1p(-jnp.exp(delta))
+    same_sign_log = big_log_safe + jnp.log1p(jnp.exp(delta))
+    opp_sign_log = big_log_safe + jnp.log1p(-jnp.exp(delta))
 
     # Default: one of the operands is zero, so result is the other.
     # (We compute `big_log` and `big_sign` which already reflect "the

tests/test_log_space.py

@@ -140,6 +140,34 @@ def test_log_add_with_zero():
     np.testing.assert_allclose(log_to_raw(res), a + b, rtol=1e-13)
 
 
+def test_log_add_both_zero_grad_finite():
+    """Both operands structural zero: forward is the zero, gradient stays finite.
+
+    Without the big_log_safe shift in log_add, ``small_log - big_log`` evaluates
+    to ``(-inf) - (-inf) = NaN`` along the dangerous (same_sign / opp_sign)
+    branches; the forward is masked correctly via ``where``, but reverse-mode
+    autodiff still differentiates through the NaN computation and propagates a
+    NaN gradient.  This regression locks in the shifted-reference behaviour.
+    """
+    from jet_array.log_space import LogSeries
+
+    zero_a = raw_to_log(jnp.array(0.0))
+    zero_b = raw_to_log(jnp.array(0.0))
+
+    res = log_add(zero_a, zero_b)
+    assert float(res.sign) == 0.0
+    assert bool(jnp.isneginf(res.log_mag))
+
+    # Differentiate the result's log_mag w.r.t. a sign-zero LogSeries whose
+    # log_mag is exactly -inf.  Pre-fix this returned NaN.
+    def loss(la):
+        a_perturbed = LogSeries(sign=zero_a.sign, log_mag=la)
+        return log_add(a_perturbed, zero_b).log_mag
+
+    g = jax.grad(loss)(zero_a.log_mag)
+    assert not bool(jnp.isnan(g)), f"gradient is NaN: {g}"
+
+
 def test_log_sum_signed_logsumexp():
     v = jnp.array([1.0, -2.0, 3.0, -4.0, 5.0])
     res = log_sum(raw_to_log(v), axis=0)