Repository: pyro-ppl/numpyro (python)

Issue:
log1m_exp and log_diff_exp functions
When writing custom distributions, it is often helpful to have numerically stable implementations of `log_diff_exp(a, b) := log(exp(a) - exp(b))` and particularly `log1m_exp(x) := log(1 - exp(x))`. The naive implementations are not stable for many probabilistic programming use cases, and so probabilistic programming languages including [Stan](https://mc-stan.org/docs/2_29/functions-reference/composed-functions.html) and [PyMC](https://docs.pymc.io/en/v3/api/math.html) provide numerically-stable implementations (typically following [Machler, 2012](https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf)) of these functions.

As far as I can tell, Numpyro does not, and they are not present in Jax. 

I wonder whether it would be worth providing them. I have written basic implementations following Machler for my own use. I would happily make a PR including them, but someone more experienced could probably write better/more idiomatic ones.

```
import jax.numpy as jnp


def log1m_exp(x):
    """
    Numerically stable calculation
    of the quantity log(1 - exp(x)),
    following the algorithm of
    Machler [1]. This is
    the algorithm used in TensorFlow Probability,
    PyMC, and Stan, but it is not provided
    yet with Numpyro.

    Currently returns NaN for x > 0,
    but may be modified in the future
    to throw a ValueError

    [1] https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf
    """
    # return 0. rather than -0. if
    # we get a negative exponent that exceeds
    # the floating point representation
    arr_x = 1.0 * jnp.array(x)
    oob = arr_x < jnp.log(jnp.finfo(
        arr_x.dtype).smallest_normal)
    mask = arr_x > -0.6931472  # appox -log(2)
    more_val = jnp.log(-jnp.expm1(arr_x))
    less_val = jnp.log1p(-jnp.exp(arr_x))

    return jnp.where(
        oob,
        0.,
        jnp.where(
            mask,
            more_val,
            less_val))


def log_diff_exp(a, b):
    # note that following Stan,
    # we want the log diff exp
    # of -inf, -inf to be -inf,
    # not nan, because that
    # corresponds to log(0 - 0) = -inf
    mask = a > b
    masktwo = (a == b) & (a < jnp.inf)
    return jnp.where(mask,
                     1.0 * a + log1m_exp(
                         1.0 * b - 1.0 * a),
                     jnp.where(masktwo,
                               -jnp.inf,
                               jnp.nan))

```


Pull request description:
Add `log1mexp` and `logdiffexp` functions
## Summary
Adds numerically stable implementations of `log1mexp` and `logdiffexp`, following [Machler 2012](https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf). Provides a custom jvp for `log1mexp` to increase numerical stability near 0.

Closes #1368 

## Changes
- Adds functions and updates documentation to include them.
- Adds tests for the functions.

## Testing / documentation.
- Verified `test_distributions_util.py` passes locally, including new tests.
- Confirmed documentation renders locally and that new docs render as expected.

## Code cleanliness and style
`pre-commit` and linting / styling pass locally. 

Patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 75363953..793c886a 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -31,7 +31,7 @@ jobs:
         python -m pip install --upgrade pip
         pip install jaxlib
         pip install jax
-        pip install .[doc,test]
+        pip install '.[doc,test]'
         pip install https://github.com/pyro-ppl/funsor/archive/master.zip
         pip install -r docs/requirements.txt
         pip freeze
@@ -70,7 +70,7 @@ jobs:
         pip install jaxlib
         pip install jax
         pip install https://github.com/pyro-ppl/funsor/archive/master.zip
-        pip install -e .[dev,test]
+        pip install -e '.[dev,test]'
         pip freeze
     - name: Test with pytest
       run: |
@@ -109,7 +109,7 @@ jobs:
         pip install jaxlib
         pip install jax
         pip install https://github.com/pyro-ppl/funsor/archive/master.zip
-        pip install -e .[dev,test]
+        pip install -e '.[dev,test]'
         pip freeze
     - name: Test with pytest
       run: |
@@ -160,7 +160,7 @@ jobs:
         pip install jaxlib
         pip install jax
         pip install https://github.com/pyro-ppl/funsor/archive/master.zip
-        pip install -e .[dev,examples,test]
+        pip install -e '.[dev,examples,test]'
         pip freeze
     - name: Test with pytest
       run: |
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index ab02f3e0..605d56b0 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -35,4 +35,4 @@ repos:
       - id: codespell
         stages: [pre-commit, commit-msg]
         args:
-          [--ignore-words-list, "Teh,aas", --check-filenames, --skip, "*.ipynb"]
+          [--ignore-words-list, "Teh,aas,ans", --check-filenames, --skip, "*.ipynb"]
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index fcfad550..0977dc2b 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -12,7 +12,7 @@ To set up local development environment, install NumPyro from source:
 ```sh
 git clone https://github.com/pyro-ppl/numpyro.git
 # install jax/jaxlib first for CUDA support
-pip install -e .[dev,test,doc,examples]  # contains additional dependencies for NumPyro development
+pip install -e '.[dev,test,doc,examples]'  # contains additional dependencies for NumPyro development
 ```
 
 For running `make doctest`, [install pandoc](https://pandoc.org/installing.html).
diff --git a/Makefile b/Makefile
index c5d6d4a9..64235143 100644
--- a/Makefile
+++ b/Makefile
@@ -14,7 +14,7 @@ format: license FORCE
 	ruff format .
 
 install: FORCE
-	pip install -e .[dev,doc,test,examples]
+	pip install -e '.[dev,doc,test,examples]'
 
 doctest: FORCE
 	JAX_PLATFORM_NAME=cpu $(MAKE) -C docs doctest
diff --git a/README.md b/README.md
index 71363610..ccdcf4be 100644
--- a/README.md
+++ b/README.md
@@ -240,13 +240,13 @@ In case of compatibility issues arise during execution of the above command, you
 compatible CPU version of JAX with
 
 ```
-pip install numpyro[cpu]
+pip install 'numpyro[cpu]'
 ```
 
 To use **NumPyro on the GPU**, you need to install CUDA first and then use the following pip command:
 
 ```
-pip install numpyro[cuda] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+pip install 'numpyro[cuda]' -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
 If you need further guidance, please have a look at the [JAX GPU installation instructions](https://github.com/jax-ml/jax#pip-installation-gpu-cuda).
@@ -265,7 +265,7 @@ You can also install NumPyro from source:
 git clone https://github.com/pyro-ppl/numpyro.git
 cd numpyro
 # install jax/jaxlib first for CUDA support
-pip install -e .[dev]  # contains additional dependencies for NumPyro development
+pip install -e '.[dev]'  # contains additional dependencies for NumPyro development
 ```
 
 You can also install NumPyro with conda:
diff --git a/docker/dev/Dockerfile b/docker/dev/Dockerfile
index da336312..409ff279 100644
--- a/docker/dev/Dockerfile
+++ b/docker/dev/Dockerfile
@@ -20,9 +20,9 @@ RUN apt update && \
 ENV PATH=/root/.local/bin:$PATH
 
 # install python packages via pip
-RUN pip3 install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+RUN pip3 install 'jax[cuda]' -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
 # clone the numpyro git repository and run pip install
 RUN git clone https://github.com/pyro-ppl/numpyro.git && \
     cd numpyro && \
-    pip3 install -e .[dev, test]  # contains additional dependencies for NumPyro development
+    pip3 install -e '.[dev, test]'  # contains additional dependencies for NumPyro development
diff --git a/docker/release/Dockerfile b/docker/release/Dockerfile
index 73755e96..abb87413 100644
--- a/docker/release/Dockerfile
+++ b/docker/release/Dockerfile
@@ -20,4 +20,4 @@ ENV PATH=/root/.local/bin:$PATH
 RUN pip3 install --user \
     # we pull wheels from google's api as per https://github.com/jax-ml/jax#installation
     # the pre-compiled wheels that google provides work for now. This may change in the future (and necessitate building from source)
-    numpyro[cuda] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+    'numpyro[cuda]' -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
diff --git a/docs/source/distributions.rst b/docs/source/distributions.rst
index ba1a8272..cbf837e1 100644
--- a/docs/source/distributions.rst
+++ b/docs/source/distributions.rst
@@ -1100,3 +1100,15 @@ BlockNeuralAutoregressiveTransform
     :undoc-members:
     :show-inheritance:
     :member-order: bysource
+
+
+Utilities
+---------
+
+log1mexp
+^^^^^^^^
+.. autofunction:: numpyro.distributions.util.log1mexp
+
+logdiffexp
+^^^^^^^^^^
+.. autofunction:: numpyro.distributions.util.logdiffexp
diff --git a/numpyro/distributions/util.py b/numpyro/distributions/util.py
index d32d2da2..ddfd02d6 100644
--- a/numpyro/distributions/util.py
+++ b/numpyro/distributions/util.py
@@ -13,6 +13,7 @@ from jax import jit, lax, random, vmap
 import jax.numpy as jnp
 from jax.scipy.linalg import solve_triangular
 from jax.scipy.special import digamma
+from jax.typing import ArrayLike
 
 from numpyro.util import not_jax_tracer
 
@@ -419,6 +420,63 @@ def logmatmulexp(x, y):
     return xy + x_shift + y_shift
 
 
+@jax.custom_jvp
+def log1mexp(x: ArrayLike) -> ArrayLike:
+    """
+    Numerically stable calculation of the quantity
+    :math:`\\log(1 - \\exp(x))`, following the algorithm
+    of `Mächler 2012`_.
+
+    .. _Mächler 2012: https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf
+
+    Returns ``-jnp.inf`` when ``x == 0`` and ``jnp.nan``
+    when ``x > 0``.
+
+    :param x: A number or array of numbers.
+    :return: The value of :math:`\\log(1 - \\exp(x))`.
+    """
+    return jnp.where(
+        x > -0.6931472,  # approx log(2)
+        jnp.log(-jnp.expm1(x)),
+        jnp.log1p(-jnp.exp(x)),
+    )
+
+
+# Custom jvp for log1mexp to handle the gradient when x is near 0.
+#
+# Inspired by the approach taken here for the function log1mexp(-x):
+# https://github.com/google-research/google-research/blob/14e984cdb8630a7e3d210dff8760fc06d490fc4b/diffusion_distillation/diffusion_distillation/utils.py#L364-L370
+# That code is (c) 2024 The Google Research Authors and licensed under
+# an Apache 2.0 License.
+log1mexp.defjvps(lambda t, ans, x: -t / jnp.expm1(-x))
+
+
+def logdiffexp(a: ArrayLike, b: ArrayLike) -> ArrayLike:
+    """
+    Numerically stable calculation of the
+    quantity :math:`\\log(\\exp(a) - \\exp(b))`,
+    provided :math:`+\\infty > a \\ge b`,
+    following the algorithm of `Mächler 2012`_.
+
+    .. _Mächler 2012: https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf
+
+    Returns ``-jnp.inf`` when ``a == b``,
+    including when ``a == b == -jnp.inf``,
+    since this corresponds to ``jnp.log(0)``.
+    Returns ``jnp.nan`` when ``a < b`` or
+    ``a == jnp.inf``.
+
+    :param a: A number or array of numbers.
+    :param b: A number or array of numbers.
+    :return: The value of :math:`\\log(\\exp(a) - \\exp(b))`.
+    """
+    return jnp.where(
+        (a < jnp.inf) & (a > b),
+        a + log1mexp(b - a),
+        jnp.where(a == b, -jnp.inf, jnp.nan),
+    )
+
+
 def clamp_probs(probs):
     finfo = jnp.finfo(jnp.result_type(probs, float))
     return jnp.clip(probs, finfo.tiny, 1.0 - finfo.eps)
diff --git a/setup.py b/setup.py
index d062fbfb..26ea1d53 100644
--- a/setup.py
+++ b/setup.py
@@ -82,7 +82,7 @@ setup(
         ],
         "cpu": f"jax[cpu]{_jax_version_constraints}",
         # TPU and CUDA installations, currently require to add package repository URL, i.e.,
-        # pip install numpyro[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html
+        # pip install 'numpyro[cuda]' -f https://storage.googleapis.com/jax-releases/jax_releases.html
         "tpu": f"jax[tpu]{_jax_version_constraints}",
         "cuda": f"jax[cuda]{_jax_version_constraints}",
     },


Test patch:
diff --git a/test/test_distributions_util.py b/test/test_distributions_util.py
index 84af13fc..14b89143 100644
--- a/test/test_distributions_util.py
+++ b/test/test_distributions_util.py
@@ -4,14 +4,15 @@
 from numbers import Number
 
 import numpy as np
-from numpy.testing import assert_allclose
+from numpy.testing import assert_allclose, assert_array_almost_equal, assert_array_equal
 import pytest
 import scipy
 
 import jax
-from jax import lax, random, vmap
+from jax import grad, lax, random, vmap
 import jax.numpy as jnp
 from jax.scipy.special import expit, xlog1py, xlogy
+from jax.test_util import check_grads
 
 import numpyro.distributions as dist
 from numpyro.distributions.util import (
@@ -20,6 +21,8 @@ from numpyro.distributions.util import (
     binomial,
     categorical,
     cholesky_update,
+    log1mexp,
+    logdiffexp,
     multinomial,
     safe_normalize,
     vec_to_tril_matrix,
@@ -78,6 +81,156 @@ def test_categorical_stats(p):
     assert_allclose(counts / float(n), p, atol=0.01)
 
 
+@pytest.mark.parametrize("x", [-80.5632, -0.32523, -0.5, -20.53, -8.032])
+def test_log1mexp_grads(x):
+    check_grads(log1mexp, (x,), order=3)
+
+
+@pytest.mark.parametrize(
+    "x, expected",
+    [
+        (jnp.array([0.01, 0, -jnp.inf]), jnp.array([jnp.nan, -jnp.inf, 0])),
+        (0.001, jnp.nan),
+        (0, -jnp.inf),
+        (-jnp.inf, 0),
+    ],
+)
+def test_log1mexp_bounds_handling(x, expected):
+    """
+    log1mexp(x) should be nan for x > 0.
+
+    log1mexp(x) should be -inf for x == 0.
+
+    log1mexp(-inf) should be 0.
+
+    This should work vectorized and not interfere
+    with other calculations.
+    """
+    assert_array_equal(log1mexp(x), expected)
+
+
+@pytest.mark.parametrize("x", [jnp.array([-0.6, -8.32, -3]), -2.5, -0.01])
+def test_log1mexp_agrees_with_basic(x):
+    """
+    log1mexp should agree with a basic implementation
+    for values where the basic implementation is stable.
+    """
+    assert_array_almost_equal(log1mexp(x), jnp.log(1 - jnp.exp(x)))
+
+
+def test_log1mexp_stable():
+    """
+    log1mexp should be stable at (negative) values of
+    x that very small and very large in absolute
+    value, where the basic implementation is not.
+    """
+
+    def basic(x):
+        return jnp.log(1 - jnp.exp(x))
+
+    # this should perhaps be made finfo-aware
+    assert jnp.isinf(basic(-1e-20))
+    assert not jnp.isinf(log1mexp(-1e-20))
+    assert_array_almost_equal(log1mexp(-1e-20), jnp.log(-jnp.expm1(-1e-20)))
+    assert abs(basic(-50)) < abs(log1mexp(-50))
+    assert_array_almost_equal(log1mexp(-50), jnp.log1p(-jnp.exp(-50)))
+
+
+@pytest.mark.parametrize("x", [-30.0, -2.53, -1e-4, -1e-9, -1e-15, -1e-40])
+def test_log1mexp_grad_stable(x):
+    """
+    Custom JVP for log1mexp should make gradient computation
+    numerically stable, even near zero, where the basic approach
+    can encounter divide-by-zero problems and yield nan.
+    The two approaches should produce almost equal answers elsewhere.
+    """
+
+    def log1mexp_no_custom(x):
+        return jnp.where(
+            x > -0.6931472,  # approx log(2)
+            jnp.log(-jnp.expm1(x)),
+            jnp.log1p(-jnp.exp(x)),
+        )
+
+    grad_custom = grad(log1mexp)(x)
+    grad_no_custom = grad(log1mexp_no_custom)(x)
+
+    assert_array_almost_equal(grad_custom, -1 / jnp.expm1(-x))
+
+    if not jnp.isnan(grad_no_custom):
+        assert_array_almost_equal(grad_custom, grad_no_custom)
+
+
+@pytest.mark.parametrize(
+    "a, b", [(-20.0, -35.0), (-0.32523, -0.34), (20.53, 19.035), (8.032, 7.032)]
+)
+def test_logdiffexp_grads(a, b):
+    check_grads(logdiffexp, (a, b), order=3, rtol=0.01)
+
+
+@pytest.mark.parametrize(
+    "a, b, expected",
+    [
+        (
+            jnp.array([jnp.inf, 0, 6.5, 4.99999, -jnp.inf]),
+            jnp.array([5, 0, 6.5, 5, -jnp.inf]),
+            jnp.array([jnp.nan, -jnp.inf, -jnp.inf, jnp.nan, -jnp.inf]),
+        ),
+        (jnp.inf, 0.3532, jnp.nan),
+        (0, 0, -jnp.inf),
+        (-jnp.inf, -jnp.inf, -jnp.inf),
+        (5.6, 5.6, -jnp.inf),
+        (1e34, 1e34 / 0.9999, jnp.nan),
+    ],
+)
+def test_logdiffexp_bounds_handling(a, b, expected):
+    """
+    Test bounds handling for logdiffexp.
+
+    logdiffexp(jnp.inf, anything) should be nan,
+
+    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))
+    """
+    assert_array_equal(logdiffexp(a, b), expected)
+
+
+@pytest.mark.parametrize(
+    "a, b", [(jnp.array([53, 23.532, 8, -1.35]), jnp.array([56, -63.2, 2, -5.32]))]
+)
+def test_logdiffexp_agrees_with_basic(a, b):
+    """
+    logdiffexp should agree with a basic implementation
+    for values at which the basic implementation is stable.
+    """
+    assert_array_almost_equal(logdiffexp(a, b), jnp.log(jnp.exp(a) - jnp.exp(b)))
+
+
+@pytest.mark.parametrize("a, b", [(500, 499), (-499, -500), (500, 500)])
+def test_logdiffexp_stable(a, b):
+    """
+    logdiffexp should be numerically stable at values
+    where the basic implementation is not.
+    """
+
+    def basic(a, b):
+        return jnp.log(jnp.exp(a) - jnp.exp(b))
+
+    if a > 0 or a == b:
+        assert jnp.isnan(basic(a, b))
+    else:
+        assert basic(a, b) == -jnp.inf
+    result = logdiffexp(a, b)
+    assert not jnp.isnan(result)
+    if not a == b:
+        assert result < a
+    else:
+        assert result == -jnp.inf
+
+
 @pytest.mark.parametrize(
     "p, shape",
     [