Fix behavior of logdiffexp(inf,inf), add test that would have caught. Fix typos (#2007)

Co-authored-by: test <[email protected]>

Author: dylanhmorris

numpyro/contrib/einstein/steinvi.py

@@ -738,7 +738,7 @@ def __init__(
     ):
         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 @@ def init(self, rng_key, num_steps, *args, **kwargs):
 
         :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`.
         """

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/test_distributions_util.py

@@ -180,20 +180,23 @@ 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):
     """
     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)