From 9f109139c48fe1ad97d05c4bc54fc450929af809 Mon Sep 17 00:00:00 2001 From: Marius1311 Date: Wed, 5 Aug 2026 10:44:21 +0200 Subject: [PATCH] Drop tensorflow_probability so scenvi imports on current jax `pip install scenvi` currently produces a package that cannot be imported at all. tensorflow_probability is pinned to ^0.22.0, i.e. <0.23, and tfp 0.22 fails against any recent jax: AttributeError: module 'jax.interpreters.xla' has no attribute 'pytype_aval_mappings' Re-pinning does not fix this for long: tfp's last release is 0.25.0 (November 2024) and it is unmaintained, while jax is at 0.11. The dependency is also far larger than what is used -- four log-densities and one matrix helper, all closed-form. So they are written out against jax directly and tfp is dropped. * scenvi/_dists.py implements Poisson, negative binomial, zero-inflated negative binomial and unit-variance normal log-densities with jax.scipy.special.gammaln / jax.nn.log_sigmoid, plus fill_triangular. * tests/test_dists.py checks every one of them against the tfp original to rtol 1e-11 in float64, including saturated logits and large counts. It skips where tfp is absent, which after this commit is everywhere except a deliberately pinned environment; it is the record that the replacements reproduce what they replaced. Two details worth flagging, both caught by that differential test: * tfp builds Inflated's mixture from categorical logits [d, -d], whose difference is 2d. The zero-inflation weight is therefore sigmoid(2d), not sigmoid(d). Using the latter silently rescales the parameter. * fill_triangular's fill order is not the obvious one: the last m - n entries are laid down first, then the whole vector reversed on top, then the lower triangle taken. Removing tfp is necessary but not sufficient, because the flax pin has gone the same way: ^0.10.4 caps below 0.11, and flax 0.10.7 calls jax.core.get_opaque_trace_state, which jax 0.11 removed. flax, optax and clu therefore lose their upper caps, and jax becomes a direct dependency rather than one inherited from flax. The bounds are written >= rather than ^ deliberately -- poetry's caret on a 0.x version caps at the next minor, which is how these pins became stale in the first place. Verified on a clean install resolving to jax 0.11.0, flax 0.12.8, optax 0.2.8, clu 0.0.12 and no tfp: `import scenvi` and `scenvi.ENVI` both work, and the suite passes. Also still passes against jax 0.4.23 with tfp 0.22 installed, where the differential tests run. Co-Authored-By: Claude Opus 5 (1M context) --- pyproject.toml | 8 +-- scenvi/ENVI.py | 4 +- scenvi/_dists.py | 82 +++++++++++++++++++++++++++--- tests/test_dists.py | 121 ++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 202 insertions(+), 13 deletions(-) create mode 100644 tests/test_dists.py diff --git a/pyproject.toml b/pyproject.toml index 9e36d82..b8da1cd 100644 --- a/pyproject.toml +++ b/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] diff --git a/scenvi/ENVI.py b/scenvi/ENVI.py index 1a92ba7..51c7f37 100644 --- a/scenvi/ENVI.py +++ b/scenvi/ENVI.py @@ -7,7 +7,6 @@ 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 @@ -15,6 +14,7 @@ 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): diff --git a/scenvi/_dists.py b/scenvi/_dists.py index 89488bf..328a431 100644 --- a/scenvi/_dists.py +++ b/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))) diff --git a/tests/test_dists.py b/tests/test_dists.py new file mode 100644 index 0000000..350e3ba --- /dev/null +++ b/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, + )