Repository: pyro-ppl/numpyro
Pull request #2007 (2025-03-12 14:02:01)

Fix behavior of `logdiffexp(jnp.inf, jnp.inf)`, add test that would have caught it.
Desired output for `logdiffexp(jnp.inf, jnp.inf)` is `jnp.nan`. We were getting `-jnp.inf` because `logdiffexp(a, a) = -jnp.inf` for any `a != jnp.inf`, including `-jnp.inf`. Added:
1. conditional to fix the issue 
2. a test of `logdiffexp` behavior at `jnp.inf, jnp.inf` that fails without these changes and passes with them.

Also fixes typo in documentation caught be codespell and clarifies an explanatory comment.


Patch:
diff --git a/numpyro/contrib/einstein/steinvi.py b/numpyro/contrib/einstein/steinvi.py
index 81fd9fd2..aedbbf94 100644
--- a/numpyro/contrib/einstein/steinvi.py
+++ b/numpyro/contrib/einstein/steinvi.py
@@ -738,7 +738,7 @@ class ASVGD(SVGD):
     ):
         assert num_cycles > 0, f"The number of cycles must be >0. Got {num_cycles}."
         assert transition_speed > 0, (
-            f"The transtion speed must be >0. Got {transition_speed}."
+            f"The transition speed must be >0. Got {transition_speed}."
         )
 
         self.num_cycles = num_cycles
@@ -785,7 +785,7 @@ class ASVGD(SVGD):
 
         :param jax.random.PRNGKey rng_key: Random number generator seed.
         :param args: Positional arguments to the model and guide.
-        :param num_steps: Totat number of steps in the optimization.
+        :param num_steps: Total number of steps in the optimization.
         :param kwargs: Keyword arguments to the model and guide.
         :return: Initial :data:`ASVGDState`.
         """
diff --git a/numpyro/distributions/util.py b/numpyro/distributions/util.py
index 24140fc9..08c9a15a 100644
--- a/numpyro/distributions/util.py
+++ b/numpyro/distributions/util.py
@@ -448,7 +448,7 @@ def log1mexp(x: ArrayLike) -> ArrayLike:
     :return: The value of :math:`\\log(1 - \\exp(x))`.
     """
     return jnp.where(
-        x > -0.6931472,  # approx log(2)
+        x > -0.6931472,  # approx -log(2)
         jnp.log(-jnp.expm1(x)),
         jnp.log1p(-jnp.exp(x)),
     )
@@ -485,7 +485,7 @@ def logdiffexp(a: ArrayLike, b: ArrayLike) -> ArrayLike:
     return jnp.where(
         (a < jnp.inf) & (a > b),
         a + log1mexp(b - a),
-        jnp.where(a == b, -jnp.inf, jnp.nan),
+        jnp.where((a < jnp.inf) & (a == b), -jnp.inf, jnp.nan),
     )
 
 


Test patch:
diff --git a/test/test_distributions_util.py b/test/test_distributions_util.py
index 2aeeaf4d..213acef6 100644
--- a/test/test_distributions_util.py
+++ b/test/test_distributions_util.py
@@ -180,7 +180,9 @@ def test_logdiffexp_grads(a, b):
         (0, 0, -np.inf),
         (-np.inf, -np.inf, -np.inf),
         (5.6, 5.6, -np.inf),
+        (1e34, 1e34, -np.inf),
         (1e34, 1e34 / 0.9999, np.nan),
+        (np.inf, np.inf, np.nan),
     ],
 )
 def test_logdiffexp_bounds_handling(a, b, expected):
@@ -188,12 +190,13 @@ def test_logdiffexp_bounds_handling(a, b, expected):
     Test bounds handling for logdiffexp.
 
     logdiffexp(jnp.inf, anything) should be nan,
+    including logdiffexp(jnp.inf, jnp.inf).
 
     logdiffexp(a, b) for a < b should be nan, even if numbers
     are very close.
 
     logdiffexp(a, b) for a == b should be -jnp.inf
-    even if a == b == -jnp.inf (log(0 - 0))
+    even if a == b == -jnp.inf (log(0 - 0)).
     """
     a = jnp.asarray(a)
     b = jnp.asarray(b)