ott-jax incompatible with JAX >= 0.9 (is_vmappable removed) — online/batched OT broken

[Marius1311]

Description

Any online / batched OT (i.e. solving or sparsifying with a PointCloud geometry that has a batch_size) currently fails with:

AttributeError: module 'jax.interpreters.batching' has no attribute 'is_vmappable'

This affects both *.solve(batch_size=...) and BaseSolverOutput.sparsify(...) on an online PointCloud. The dense/offline path (no batch_size) still works. Large-data users hit this because they set batch_size in solve to scale (see scverse/cellrank#1146).

Reproduced in a fresh dev env (jax==0.10.1, ott-jax==0.6.0, Python 3.14):

import jax.numpy as jnp
from ott.geometry.pointcloud import PointCloud
from ott.problems.linear.linear_problem import LinearProblem
from ott.solvers.linear.sinkhorn import Sinkhorn

g = PointCloud(jnp.zeros((4, 3)), jnp.zeros((5, 3)), epsilon=0.1, batch_size=2)  # online
Sinkhorn()(LinearProblem(g))   # -> AttributeError: ... has no attribute 'is_vmappable'

Root cause

ott's batched_vmap (used by online PointCloud.apply_lse_kernel) relies on jax.interpreters.batching.is_vmappable. JAX kept that symbol (deprecated) through 0.8.2 and removed it in 0.9.0. Since moscot pins jax>=0.6.1 with no upper bound, a fresh install resolves to jax ≥ 0.9 and breaks.

Upstream status (important nuance)

OTT main attempted a fix in ott-jax/ott#673 by switching the import from the public jax.interpreters.batching to the private jax._src.interpreters.batching. However, this fix is ineffective on jax ≥ 0.9: from jax._src.interpreters import batching resolves the package attribute jax._src.interpreters.batching, which jax re-exports as the public module (the one missing is_vmappable) — even though sys.modules['jax._src.interpreters.batching'] is the real private module that still has the symbol. So installing ott-jax from git main does not resolve the breakage on jax 0.10.1 (verified). A robust upstream fix would force the submodule (e.g. import jax._src.interpreters.batching as batching) or drop the is_vmappable dependency; worth tracking / filing upstream.

Note: the latest ott-jax release (0.6.0, 2025-11-04) predates the jax 0.9 removal and also uses the public import.

Suggested fix

• Near-term: cap jax < 0.9 in pyproject.toml so online OT works with the current ott-jax release.
• Longer-term: once a fixed ott-jax release is available, bump ott-jax and drop the cap.

Links

• Surfaces for users in RealTimeKernel.from_moscot(tp) Out of Memory scverse/cellrank#1146 (RealTimeKernel.from_moscot OOM / failure).
• Unblocks the memory-bounded sparsify work for change batch_size in sparsify #639.

[Marius1311]

Filed upstream: ott-jax/ott#700 (with a minimal repro and a verified one-line fix). It also documents why ott-jax/ott#673 does not actually resolve this on JAX ≥ 0.9.