fix(sparsify): materialize transport-matrix rows instead of pushing indicator columns

[Marius1311]

Problem

sparsify extracts rows of the transport matrix by pushing dense identity blocks:

rows = np.asarray(self.push(np.eye(n, cols, -batch), scale_by_marginals=False)).T

ott's log-sum-exp apply materializes an [n, m, k] tensor for a k-column input, so peak memory is n · m · batch_size, not batch_size as the docstring promises. Measured on main (moscot 0.5.1, ott-jax 0.6.0 and main, jax 0.11.1), for a TemporalProblem whose dense couplings are 18 MB at n = m = 1500:

sparsify_kwargs={"batch_size": …}	8	64	256	1024 (default)
peak RSS	0.7 GB	1.7 GB	5.8 GB	17.8 GB

Isolated at the ott level, output.apply on an [n, k] input peaks at ≈ 3 · n · m · k · 4 bytes. At n = m = 8000 the default batch_size would need ~262 GB to sparsify a 256 MB plan. The knob you have to shrink to survive is also the one that makes the result worse (below), and #843's row-oriented loop did not change this — it fixed the orientation but kept the identity push.

This is what large-data users hit downstream: scverse/cellrank#1146.

Fix

Rows are built directly, behind one seam — BaseDiscreteSolverOutput._row_materializer(), a factory returning a callable over row indices (built once per sparsify call, so per-output work such as resolving a geometry is hoisted out of the block loop):

• low-rank (LRSinkhornOutput, LRGWOutput): T[rows] = (Q diag(1/g) Rᵀ)[rows], exactly ott's own matrix. Never touches a geometry — which matters, since LRGWOutput.geom re-linearizes on every access.
• entropic (SinkhornOutput, and GWOutput via linear_state): T[rows] = exp((f_i + g_j − C_ij) / ε), with the cost rows from cost_fn.all_pairs(x[rows], y) (PointCloud), cost_1[rows] @ cost_2ᵀ + bias (LRCGeometry) or cost_matrix[rows] (dense Geometry).
• MatrixSolverOutput: slices what it already holds.
• anything else — notably GraphOTTOutput, whose geometry is a Geodesic and whose rows are a sub-block of an expanded (n+m)² problem — keeps the push-based default, which is cheap there: applying a sparse graph kernel costs nnz per column, not n · m.

Two correctness traps this had to avoid, both verified:

• geom.subset(row_ixs=…) cannot be used. scale_cost and relative_epsilon travel as unevaluated strings in aux_data, so a subset geometry re-derives inv_scale_cost and epsilon from the subset. With scale_cost="mean" (what moscot solves with) that shifted inv_scale_cost 0.05235 → 0.05396 and produced 15% value error, silently. The materializer captures the parent's resolved values instead. LRCGeometry.subset is worse — the inherited Geometry.subset slices the two factors as if they were cost/kernel.
• GWOutput._rescale_factor (sqrt(old_transport_mass / linear_state.transport_mass)) must be applied, as matrix and apply both do. It is ~1.0 on converged balanced problems, so it hides easily; there is a test that forces it to 2.0.

Results

Peak memory at n = m = 1500 drops from 17.8 GB to 84 MB (212×) and no longer scales with batch_size (ratio across 64 → 1024 goes from 4.8× to 0.6×); min_row is now cheaper than materializing the dense plan. Reconstruction is bit-exact (rel_err = 0.0) for Sinkhorn (offline / online / low-rank-cost geometry), low-rank Sinkhorn, GW, low-rank GW and FGW, across scale_cost ∈ {1.0, "mean", "max_cost"}.

Behavioural changes

min_row is now exact and independent of batch_size. The intent — the largest threshold that keeps every row, min_i max_j T_ij, so the result can still be normalized into a Markov chain — was computed over column batches, which can only underestimate it. At n = m = 1500 it kept 182,645 entries at batch_size=1024 and 1,227,714 at 64, against 2,319 for the exact rule; the threshold and the values also came from two different float paths (pull vs push), which could empty a row outright (I measured 199/200 rows non-empty at batch_size = m). Results are unchanged whenever m <= batch_size, which is why small problems never showed this; users with m > batch_size will see much sparser couplings.

Rows carrying no mass are handled centrally, for every mode: they keep no entries and take no part in choosing a threshold. No threshold can make them non-empty, while letting them participate drags any row-based threshold to 0 and disables sparsification altogether (a single massless row took a test matrix from 21% to 88% dense). A warning is emitted, since the result then cannot be normalized into a Markov chain. mass mode already behaved this way; the modes are now consistent.

percentile estimates its threshold from sampled rows rather than pulled column indicators, so it uses the same bounded path as everything else.

OTTOutput._with_batch_size is removed. It only ever rebatched a Sinkhorn output over an online PointCloud — exactly the case now built from potentials — so it became dead code.

Tests

• tests/solvers/test_base_solver.py: min_row independence from batch_size, the every-row-non-empty contract, exactness against min_i max_j T_ij, percentile row sampling, metadata propagation, a TestMinRowStructure class checking the pull proxy on a transport-plan-like matrix, and a TestMasslessRows class covering all four modes.
• tests/backends/ott/test_backend.py: TestSparsifyRows — reconstruction across 7 output flavours × 3 scale_cost × 3 batch_size, a monkeypatched push/pull/transport_matrix that must never be called, the min_row contract on real outputs, a parent-scale_cost check, and the forced GW rescale factor.
• tests/backends/ott/test_sparsify_memory.py: strict peak-RSS tests, each in a fresh subprocess (peak RSS is process-global and xdist would confound it). Marked memory and deselected by default (addopts = -m 'not memory'), since they are environment-sensitive and deliberately large enough that the old implementation exhausts memory; run them with pytest -m memory.

The new tests are discriminating: 37 fail against the current main sources and all pass with the fix.

One pre-existing test changed: test_sparsify_minrow's corrcoef(pull_sparse, pull_dense) > 0.5 proxy encoded the old, denser output — on a uniform random fixture the exact rule keeps ~1 entry per row, which cannot correlate with the dense pull. It is replaced by an exactness assertion, and the proxy is retained in TestMinRowStructure on a peaked, plan-like matrix where it is meaningful.

Verification

Full suite: 839 passed. The 27 failures in test_annotation_mapping / test_set_graph_xy are pre-existing in my environment (ott-jax main + jax 0.11) — identical counts with and without this change. black/isort/ruff clean. End-to-end from cellrank (RealTimeKernel.from_moscot(..., sparse_mode="min_row"), n = m = 1500): sparse couplings, no empty rows, row-stochastic transition matrix, and its test_from_moscot tests pass.

Out of scope, filed separately: GraphOTTOutput overrides shape but inherits transport_matrix from OTTOutput, returning the full expanded (n+m)² matrix.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 80.00000% with 16 lines in your changes missing coverage. Please review.
✅ Project coverage is 76.54%. Comparing base (1d4a4bc) to head (4341072).
⚠️ Report is 8 commits behind head on main.

Files with missing lines	Patch %	Lines
src/moscot/base/output.py	75.67%	8 Missing and 1 partial ⚠️
src/moscot/backends/ott/output.py	83.72%	4 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #846      +/-   ##
==========================================
- Coverage   76.62%   76.54%   -0.08%     
==========================================
  Files          36       36              
  Lines        4175     4217      +42     
  Branches      670      677       +7     
==========================================
+ Hits         3199     3228      +29     
- Misses        679      689      +10     
- Partials      297      300       +3     

Files with missing lines	Coverage Δ
src/moscot/backends/ott/output.py	81.87% <83.72%> (-0.68%)	⬇️
src/moscot/base/output.py	81.49% <75.67%> (-2.08%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[selmanozleyen]

why didn't the first fix work I wonder and I wonder if there are any regression before and after #843 . I'd basically try to test these in all cases to see what regresses and what doesn't

[Marius1311]

why didn't the first fix work I wonder and I wonder if there are any regression before and after #843 . I'd basically try to test these in all cases to see what regresses and what doesn't

There's a new test group, pytest.mark.memory, that's excluded in CI. If you run these locally, you should be able to reproduce that #843 didn't fix it.

[selmanozleyen]

why use this jax pattern in numpy? Like returning a function to call it later seems redundant no?

[selmanozleyen]

OK could you perhaps post me some more tables on each branch proving what branch consumes how much memory and time in certain configurations? I'd like to get this right this time.

Ok so here the core bottleneck from what I understand is conversion from gpu to cpu without blowing up. And because we move the gpu arrays from jax via converting to numpy arrays we can't have nicely compiled for loops hence the python for loops become painfully slow.

Can't we use jax right until the very very end? So the code path is compiled? Like is it possible to move the arrays to cpu jax and have a list of them and convert them once in the end to numpy?

[Marius1311]

Why #843 didn't fix it. It changed the loop orientation and added mass, but rows were still extracted by pushing identity blocks through the output, and that's where the memory goes: ott's LSE apply vmaps over the vector batch, so k columns cost an [n, m, k] tensor. Peak stays at n · m · batch_size.

GPU (RTX 4090, 24 GB, min_row, JAX allocator peak above the post-solve baseline, preallocation off):

n	batch_size	pre-#843 (bacec6c7)	#843 (main)	#846 (this PR)
2000	64	2,914 MB / 5.1 s	2,914 MB / 6.3 s	0 MB / 2.2 s
2000	1024 (default)	OOM	OOM	0 MB / 2.0 s
4000	64	11,489 MB / 11.8 s	11,489 MB / 12.7 s	0 MB / 2.8 s
4000	1024 (default)	OOM	OOM	0 MB / 2.2 s

"0 MB" means sparsify never pushes the allocator above what the solve already reserved, i.e. the [batch_size, m] blocks fit under that ceiling. Host side it's 33-70 MB. Note the OOMs are at the default batch_size: on a 24 GB card neither older version can sparsify a 2000x2000 plan out of the box, which is the failure in scverse/cellrank#1146.

CPU, same three versions, peak RSS above baseline / time / nnz:

n	batch_size	pre-#843	#843	#846
1000	64	753 MB / 11.7 s / 533,556	824 MB / 8.2 s / 533,556	19 MB / 1.5 s / 22,588
1000	1024	7,690 MB / 10.9 s / 22,588	7,709 MB / 6.9 s / 22,588	15 MB / 0.4 s / 22,588
2000	64	2,916 MB / 50.6 s / 2,571,718	2,932 MB / 35.3 s / 2,571,718	2 MB / 1.6 s / 98,812
2000	1024	~49 GB, not run	~49 GB, not run	11 MB / 1.5 s / 98,812

Regressions. pre-#843 and #843 agree to the entry everywhere, on both devices, so #843 neither helped nor hurt the memory. mass (added in #843) is bit-identical between #843 and #846: 192,362 nnz at n=1000, 764,584 at n=2000. The one thing that changes is min_row, and the table shows why: 533,556 entries at batch_size=64 versus 22,588 at 1024 for the same problem, in both older versions, because the threshold was computed over column batches. #846 gives 22,588 either way, which is the exact min_i max_j T_ij rule. Identical whenever m <= batch_size, so small problems never showed it. Low-rank was never affected, since LR apply is a matmul rather than an LSE kernel.

On ott. Part of this is upstream: the same call with lse_mode=False peaks at 20 MB instead of 1,314 MB (n=m=1500, k=64) and agrees to 9e-7, so the k factor is an artefact of the LSE path, not intrinsic. Happy to file it. But apply is the wrong tool for materialising rows either way, since on a dense geometry you still touch [n, m] per call, and for low-rank the rows are just Q diag(1/g) R^T.

On staying in jax. I profiled it: 83-84% of the time is jax row construction plus the device to host transfer, 16-17% is the numpy thresholding and CSR assembly. The Python loop isn't the bottleneck, 32 blocks takes 1.78 s and 2 blocks takes 1.71 s, so dropping 30 iterations buys ~4%. Everything up to the block is already compiled jax; only the [b, m] block crosses to host. Keeping blocks on device until the end can't work for the memory goal, since accumulating them is exactly the dense matrix we're avoiding, and the CSR has to be built host-side. The promising version of your idea is thresholding on device and transferring only surviving values and indices, which on GPU would move ~2,300 numbers instead of 4M, but it needs a static nnz bound to stay jittable, so I'd do it as a follow-up.

Two footnotes: GPU and CPU nnz differ marginally (98,784 vs 98,812) from float ordering, and the ~49 GB CPU cells are extrapolated from the measured scaling rather than run, the GPU OOMs being the measured version of those.

🤖 Generated with Claude Code

[selmanozleyen]

Ok perfect, I also ran them myself it all seems to make sense. But I'd like to make some changes for the code cleanliness a bit.

[Marius1311]

Sure, feel free to modify as you like.

[selmanozleyen]

Thanks a lot! This seems like a conclusive solution to it. I will make a release now