Make JAX scGen reproduce PyTorch scGen

[Zethson]

Summary

pertpy's JAX scGen has never faithfully reproduced the canonical PyTorch scGen (theislab/scgen). Diffing the JAX module against the canonical source turned up three concrete divergences, all fixed here:

1. Reconstruction loss reduced over the wrong axis (the dominant bug).

   # before — sums over the ENTIRE minibatch → scalar
   return jnp.sum((x - px) ** 2)
   # after — per cell, sum over genes → vector [batch], matching scGen's .sum(dim=1)
   return jnp.sum((x - px) ** 2, axis=-1)

   With loss = mean(0.5 * rl + 0.5 * kl * kl_weight), the old scalar rl scaled with batch_size, dwarfing the KL term and making the objective batch-size dependent.

2. Decoder was missing batch norm. Canonical DecoderSCGEN uses scvi FCLayers (default use_batch_norm=True) and SCGENVAE runs with use_batch_norm="both". pertpy built FlaxDecoder without passing use_batch_norm, so it defaulted to False. Now the decoder honors use_batch_norm/use_layer_norm like the encoder.

3. Encoder did not add var_eps. scvi's Encoder adds var_eps=1e-4 to the latent variance; the Flax encoder now does too.

The architecture otherwise already matched (n_hidden=800, n_latent=100, n_layers=2, dropout=0.2, LeakyReLU, BN momentum/eps).

Verification

• Offline numerical parity: with identical parameters and inputs, the module's training loss now equals the canonical scGen objective recomputed by hand — 6.5323769911546… vs 6.5323769911546…, matching to ~15 significant digits.
• The end-to-end train → batch_removal → predict → reg-mean/var pipeline still runs.

Tests

Added three regression tests (no new dependency) pinning each fix:

• test_scgen_reconstruction_loss_is_per_cell — reconstruction loss has shape (n_cells,).
• test_scgen_decoder_uses_batch_norm — the decoder has BatchNorm params under use_batch_norm="both".
• test_scgen_encoder_applies_var_eps — the encoder adds var_eps to the variance.

Note

This changes outputs for existing users — intentionally, since the prior behavior did not reproduce scGen.

🤖 Generated with Claude Code

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.00%. Comparing base (473da92) to head (188da9e).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1028      +/-   ##
==========================================
+ Coverage   77.95%   78.00%   +0.05%     
==========================================
  Files          50       50              
  Lines        6587     6557      -30     
==========================================
- Hits         5135     5115      -20     
+ Misses       1452     1442      -10     

Files with missing lines	Coverage Δ
pertpy/tools/_scgen/_base_components.py	96.15% <100.00%> (+2.03%)	⬆️
pertpy/tools/_scgen/_scgenvae.py	93.93% <100.00%> (+0.18%)	⬆️

... and 4 files with indirect coverage changes

🚀 New features to boost your workflow:

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