[Marius1311]

sparsify(mode="mass", max_k=...) copied every [batch_size, m] block of rows to the host to keep at most max_k entries per row. For ott outputs the rows are built on the device, so this selects there and copies only [batch_size, max_k] per block.

• BaseDiscreteSolverOutput._mass_select is a new hook for the mass rule; its default is the existing numpy selection. OTTOutput overrides it when the rows are a jax.Array and max_k is set, using jax.lax.top_k over 256 rows at a time: one top_k over a 2048 × 92,507 block needed several times the block's memory and ran 3x slower.
• It keeps the same entries as the host selection, except among tied entries (e.g. rows of low-rank outputs), where an equal value may be kept at another column. The new test compares the kept values per row.

Two Sinkhorn couplings from a single-cell time course, value=0.95, max_k=100, batch_size=2048, RTX 4090:

coupling	host selection (as in #850)	device selection
47,842 × 55,104	193 s	0.98 s
55,104 × 92,507	384 s	1.86 s

Same number of entries in both; 22 and 60 rows keep a tied entry at another column (values differ by < 1e-7). On CPU, one 2048 × 92,507 block takes 0.99 s with numpy and 0.68 s with jax.

🤖 Generated with Claude Code

https://claude.ai/code/session_01CEkHnYavwEf7n32j7rNkPv

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 76.74%. Comparing base (4b90f31) to head (4d2b5f6).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #851      +/-   ##
==========================================
+ Coverage   76.58%   76.74%   +0.15%     
==========================================
  Files          36       36              
  Lines        4224     4252      +28     
  Branches      678      679       +1     
==========================================
+ Hits         3235     3263      +28     
  Misses        689      689              
  Partials      300      300              

Files with missing lines	Coverage Δ
src/moscot/backends/ott/output.py	84.02% <100.00%> (+2.14%)	⬆️
src/moscot/base/output.py	82.42% <100.00%> (+0.37%)	⬆️

🚀 New features to boost your workflow:

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