feat: Add JointEmbeddingSCVI model

[vals]

Hi!

I curated the code from "Improving SCVI for low-count cells through self-supervised augmentation" into a pull request to add the new model class in scvi-tools.

I am working on a tutorial (that is a bit more visual and intuitive than the metrics I used in the paper). I'll PR that to the tutorials repo in a bit.

The summary below was generated by GitHub, and summarizes it nicely:

Add JointEmbeddingSCVI, a self-supervised SCVI variant that augments the standard ELBO with a cross-correlation objective (CCO) between the embedding of the data and a binomially-thinned view of it, encouraging embeddings that are robust to count dropout/noise. A VICReg-style variance term guards against dimension collapse.

The joint-embedding loss is active only during training; evaluation (get_elbo / get_marginal_ll / get_reconstruction_error) falls back to the standard reweighted ELBO, and so does any minibatch with fewer than two cells.

Contents:

• src/scvi/module/_joint_embedding_utils.py: binomial thinning, per-cell thinning-probability sampling, cross-correlation and variance losses
• src/scvi/module/_joint_embedding_vae.py: JointEmbeddingVAE module
• src/scvi/model/_joint_embedding_scvi.py: JointEmbeddingSCVI model
• tests/model/test_joint_embedding_scvi.py, tests/module/test_joint_embedding_utils.py
• docs/api (user + developer) and CHANGELOG entries

[codecov]

Codecov Report

❌ Patch coverage is 98.94737% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 88.50%. Comparing base (f34d577) to head (836fed1).

Files with missing lines	Patch %	Lines
src/scvi/external/joint_embedding_scvi/_module.py	98.03%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3883      +/-   ##
==========================================
+ Coverage   88.46%   88.50%   +0.04%     
==========================================
  Files         227      231       +4     
  Lines       22955    23050      +95     
==========================================
+ Hits        20306    20401      +95     
  Misses       2649     2649              

Flag	Coverage Δ
integration	74.19% <98.94%> (+0.10%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/scvi/external/__init__.py	100.00% <100.00%> (ø)
src/scvi/external/joint_embedding_scvi/__init__.py	100.00% <100.00%> (ø)
src/scvi/external/joint_embedding_scvi/_model.py	100.00% <100.00%> (ø)
src/scvi/external/joint_embedding_scvi/_utils.py	100.00% <100.00%> (ø)
src/scvi/external/joint_embedding_scvi/_module.py	98.03% <98.03%> (ø)

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

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

[ori-kron-wis]

close #3734

[ori-kron-wis]

Hey @vals, Thank you so much for this!

I want to make several changes and have made my own copy of this branch here: https://github.com/scverse/scvi-tools/tree/Ori_joint-embedding-scvi (I can't push to your repo), main changes:

1. The model, module, and utils files were moved to src/scvi/external/joint_embedding_scvi/ as this is an external model, and all relevant paths and indexes were updated. Tests are also under external review now.
2. Updated changelog, reference, model doc, and tutorial links (need to be added as said)
3. Verifying not using gene_likelihood == "normal" when doing use_joint_embedding

Feel free to merge that branch into this one (of course, the tutorial will also need to adapt), or continue there.