Drop tensorflow_probability so scenvi imports on current jax

pyproject.toml

@@ -9,10 +9,10 @@ readme = "README.md"
 [tool.poetry.dependencies]
 scanpy = "^1.11.0"
 python = "^3.9"
-flax = "^0.10.4"
-optax = "^0.2.4"
-tensorflow_probability = "^0.22.0"
-clu = "^0.0.11"
+jax = ">=0.4.27"
+flax = ">=0.10.4"
+optax = ">=0.2.4"
+clu = ">=0.0.11"
 tqdm = "^4.66.1"
 
 [build-system]

scenvi/ENVI.py

@@ -7,14 +7,14 @@
 import pandas as pd
 import scanpy as sc
 import sklearn.neighbors
-import tensorflow_probability.substrates.jax as jax_prob # type: ignore
 from flax import linen as nn
 from jax import jit, random
 from tqdm import trange, tqdm
 
 from scenvi._dists import (
     KL,
     AOT_Distance,
+    fill_triangular,
     log_nb_pdf,
     log_normal_pdf,
     log_pos_pdf,
@@ -368,7 +368,7 @@ def grammian_cov(self, dec_cov):
         :meta private:
         """
 
-        dec_cov = jax_prob.math.fill_triangular(dec_cov)
+        dec_cov = fill_triangular(dec_cov)
         return jnp.matmul(dec_cov, dec_cov.transpose([0, 2, 1]))
 
     def create_train_state(self, key=random.key(0), init_lr=3e-4, decay_steps=100):

scenvi/_dists.py

@@ -1,5 +1,25 @@
+"""Log-densities and matrix helpers for ENVI's decoders.
+
+These were tensorflow_probability's jax substrate. tfp's last release is 0.25.0
+(November 2024) and it is unmaintained, while its pinned version here (<0.23)
+fails to import against any recent jax:
+
+    AttributeError: module 'jax.interpreters.xla' has no attribute
+    'pytype_aval_mappings'
+
+The surface actually used was five functions, all of them closed-form, so they
+are written out here against jax directly. Each one is checked against the tfp
+original in tests/test_dists.py.
+"""
+
+import math
+
 import jax.numpy as jnp
-import tensorflow_probability.substrates.jax.distributions as jnd
+from jax.nn import log_sigmoid
+from jax.scipy.special import gammaln, xlogy
+
+#: ``log(2 * pi)``, the Gaussian normalization constant.
+LOG_TWO_PI = math.log(2.0 * math.pi)
 
 
 def KL(mean, log_std):
@@ -15,7 +35,9 @@ def log_pos_pdf(sample, l):  # noqa: E741
     :meta private:
     """
 
-    log_prob = jnd.Poisson(rate=l).log_prob(sample)
+    # Poisson(rate=l): k log(l) - l - log(k!). xlogy keeps k = 0, l = 0 at 0
+    # rather than nan, which is what tfp does.
+    log_prob = xlogy(sample, l) - l - gammaln(sample + 1.0)
     return jnp.mean(log_prob, axis=-1)
 
 
@@ -24,7 +46,11 @@ def log_nb_pdf(sample, r, p):
     :meta private:
     """
 
-    log_prob = jnd.NegativeBinomial(total_count=r, logits=p).log_prob(sample)
+    # NegativeBinomial(total_count=r, logits=p), i.e. p is the logit of the
+    # per-trial success probability:
+    #   log C(k + r - 1, k) + k log(sigmoid(p)) + r log(1 - sigmoid(p))
+    binomial_coefficient = gammaln(sample + r) - gammaln(sample + 1.0) - gammaln(r)
+    log_prob = binomial_coefficient + sample * log_sigmoid(p) + r * log_sigmoid(-p)
     return jnp.mean(log_prob, axis=-1)
 
 
@@ -33,9 +59,23 @@ def log_zinb_pdf(sample, r, p, d):
     :meta private:
     """
 
-    log_prob = jnd.Inflated(
-        jnd.NegativeBinomial(total_count=r, logits=p), inflated_loc_logits=d
-    ).log_prob(sample)
+    # Zero-inflated negative binomial: with probability sigmoid(2d) the sample is
+    # a structural zero, otherwise it is drawn from the negative binomial. Only
+    # the k = 0 branch mixes the two, and it is summed in log space.
+    #
+    # The factor of two is not a typo. tfp built the mixture's categorical from
+    # logits [d, -d], whose difference is 2d, so the inflation weight is
+    # sigmoid(2d) and not sigmoid(d). Getting this wrong rescales the
+    # zero-inflation parameter silently -- see tests/test_dists.py.
+    binomial_coefficient = gammaln(sample + r) - gammaln(sample + 1.0) - gammaln(r)
+    log_nb = binomial_coefficient + sample * log_sigmoid(p) + r * log_sigmoid(-p)
+
+    log_nb_at_zero = r * log_sigmoid(-p)
+    log_prob = jnp.where(
+        sample == 0,
+        jnp.logaddexp(log_sigmoid(2.0 * d), log_sigmoid(-2.0 * d) + log_nb_at_zero),
+        log_sigmoid(-2.0 * d) + log_nb,
+    )
     return jnp.mean(log_prob, axis=-1)
 
 
@@ -44,7 +84,7 @@ def log_normal_pdf(sample, mean):
     :meta private:
     """
 
-    log_prob = jnd.Normal(loc=mean, scale=1).log_prob(sample)
+    log_prob = -0.5 * (jnp.square(sample - mean) + LOG_TWO_PI)
     return jnp.mean(log_prob, axis=-1)
 
 
@@ -57,3 +97,31 @@ def AOT_Distance(sample, mean):
     mean = jnp.reshape(mean, [mean.shape[0], -1])
     log_prob = -jnp.square(sample - mean)
     return jnp.mean(log_prob, axis=-1)
+
+
+def fill_triangular(x):
+    """Pack ``x`` into the lower triangle of a square matrix, tfp's way.
+
+    Reproduces ``tfp.math.fill_triangular``, whose fill order is not the obvious
+    one -- the last ``m - n`` entries are laid down first, then the whole vector
+    reversed on top, and the lower triangle taken::
+
+        [1, 2, 3, 4, 5, 6] -> [[4, 0, 0],
+                               [6, 5, 0],
+                               [3, 2, 1]]
+
+    :param x: (array) ``(..., n * (n + 1) / 2)`` values to pack
+
+    :return: (array) ``(..., n, n)`` lower-triangular matrices
+
+    :meta private:
+    """
+
+    m = x.shape[-1]
+    # m = n (n + 1) / 2, so n is the positive root of n^2 + n - 2m.
+    n = int(round((math.sqrt(1.0 + 8.0 * m) - 1.0) / 2.0))
+    if n * (n + 1) // 2 != m:
+        raise ValueError(f"last dimension {m} is not a triangular number")
+
+    packed = jnp.concatenate([x[..., n:], jnp.flip(x, axis=-1)], axis=-1)
+    return jnp.tril(jnp.reshape(packed, (*x.shape[:-1], n, n)))

tests/test_dists.py

@@ -0,0 +1,121 @@
+"""The tfp-free log-densities must agree with the tensorflow_probability originals.
+
+scenvi no longer depends on tensorflow_probability, so these run only where it is
+still installed and importable -- which, given the <0.23 pin fails against any
+recent jax, in practice means a deliberately pinned environment. They are the
+record that the replacements in `scenvi/_dists.py` reproduce what they replaced.
+"""
+
+import numpy as np
+import pytest
+
+jnp = pytest.importorskip("jax.numpy")
+tfp = pytest.importorskip("tensorflow_probability.substrates.jax")
+jnd = tfp.distributions
+
+from scenvi._dists import (  # noqa: E402
+    fill_triangular,
+    log_nb_pdf,
+    log_normal_pdf,
+    log_pos_pdf,
+    log_zinb_pdf,
+)
+
+SHAPE = (16, 32)
+
+
+@pytest.fixture(autouse=True)
+def double_precision():
+    """Compare the formulae, not float32 rounding, whatever the ambient config."""
+    from jax.experimental import enable_x64
+
+    with enable_x64():
+        yield
+
+
+@pytest.fixture
+def rng():
+    return np.random.default_rng(0)
+
+
+@pytest.fixture
+def counts(rng):
+    """Non-negative counts, deliberately including exact zeros."""
+    sample = rng.poisson(2.0, size=SHAPE).astype(np.float64)
+    assert (sample == 0).any(), "the zero branch must be exercised"
+    return jnp.asarray(sample)
+
+
+def test_poisson_matches_tfp(rng, counts):
+    rate = jnp.asarray(rng.uniform(0.05, 10.0, size=SHAPE))
+
+    expected = jnp.mean(jnd.Poisson(rate=rate).log_prob(counts), axis=-1)
+    np.testing.assert_allclose(log_pos_pdf(counts, rate), expected, rtol=1e-12, atol=1e-12)
+
+
+def test_negative_binomial_matches_tfp(rng, counts):
+    total_count = jnp.asarray(rng.uniform(0.5, 20.0, size=SHAPE))
+    logits = jnp.asarray(rng.normal(0.0, 3.0, size=SHAPE))
+
+    expected = jnp.mean(
+        jnd.NegativeBinomial(total_count=total_count, logits=logits).log_prob(counts), axis=-1
+    )
+    np.testing.assert_allclose(
+        log_nb_pdf(counts, total_count, logits), expected, rtol=1e-11, atol=1e-11
+    )
+
+
+def test_zero_inflated_negative_binomial_matches_tfp(rng, counts):
+    total_count = jnp.asarray(rng.uniform(0.5, 20.0, size=SHAPE))
+    logits = jnp.asarray(rng.normal(0.0, 3.0, size=SHAPE))
+    inflation = jnp.asarray(rng.normal(0.0, 3.0, size=SHAPE))
+
+    expected = jnp.mean(
+        jnd.Inflated(
+            jnd.NegativeBinomial(total_count=total_count, logits=logits),
+            inflated_loc_logits=inflation,
+        ).log_prob(counts),
+        axis=-1,
+    )
+    np.testing.assert_allclose(
+        log_zinb_pdf(counts, total_count, logits, inflation), expected, rtol=1e-11, atol=1e-11
+    )
+
+
+def test_normal_matches_tfp(rng):
+    sample = jnp.asarray(rng.normal(size=SHAPE))
+    mean = jnp.asarray(rng.normal(size=SHAPE))
+
+    expected = jnp.mean(jnd.Normal(loc=mean, scale=1).log_prob(sample), axis=-1)
+    np.testing.assert_allclose(log_normal_pdf(sample, mean), expected, rtol=1e-12, atol=1e-12)
+
+
+@pytest.mark.parametrize("n", [1, 2, 3, 4, 8, 16])
+def test_fill_triangular_matches_tfp(rng, n):
+    x = jnp.asarray(rng.normal(size=(5, n * (n + 1) // 2)))
+
+    np.testing.assert_array_equal(fill_triangular(x), tfp.math.fill_triangular(x))
+
+
+def test_extreme_parameters_match_tfp(rng):
+    """Saturated logits and large counts, where a naive log(sigmoid(p)) underflows."""
+    sample = jnp.asarray(rng.integers(0, 500, size=SHAPE).astype(np.float64))
+    total_count = jnp.asarray(rng.uniform(1.0, 500.0, size=SHAPE))
+    logits = jnp.asarray(rng.uniform(-40.0, 40.0, size=SHAPE))
+    inflation = jnp.asarray(rng.uniform(-40.0, 40.0, size=SHAPE))
+
+    nb = jnd.NegativeBinomial(total_count=total_count, logits=logits)
+    np.testing.assert_allclose(
+        log_nb_pdf(sample, total_count, logits),
+        jnp.mean(nb.log_prob(sample), axis=-1),
+        rtol=1e-11,
+        atol=1e-11,
+    )
+    np.testing.assert_allclose(
+        log_zinb_pdf(sample, total_count, logits, inflation),
+        jnp.mean(
+            jnd.Inflated(nb, inflated_loc_logits=inflation).log_prob(sample), axis=-1
+        ),
+        rtol=1e-11,
+        atol=1e-11,
+    )