Add semi-supervised losses and test

[surajyadav-research]

#1550

fixmatch_loss — tests

• Random batched matches reference: compares output to a float32 reference implementation (hard/soft supervised labels), across different B/U/C, confidence_threshold, lambda_u, and dtype (incl. bfloat16); also checks output is finite.
• vmap correctness: jax.vmap(fixmatch_loss) matches lax.map (per-item loop) output; also checks finiteness.
• Permutation invariance: shuffling labeled batch order and unlabeled batch order doesn’t change the loss (with threshold set so ordering shouldn’t matter); checks finiteness.
• Numerical stability (extreme logits): loss stays finite for very large-magnitude logits; gradients w.r.t. labeled logits and strong unlabeled logits are finite.
• lambda_u = 0 supervised-only: returns exactly supervised cross-entropy when unsupervised weight is zero.
• Confidence threshold edges:
  • too high (e.g., >1) → no pseudo-labels used → supervised-only
  • 0.0 → unsupervised term included (loss ≥ supervised loss)
• Empty unlabeled batch: if U=0, behaves as supervised-only; finite.
• Gradient flows through strong logits: gradient w.r.t. strong logits (us) is non-zero and finite.
• bfloat16 run: smoke test that it runs in bfloat16 and returns finite output.

---

mixmatch_loss — tests

• Random batched matches reference: compares output to a float32 reference implementation (hard/soft supervised labels), different B/U/C, lambda_u, and dtype (incl. bfloat16); checks finiteness.
• vmap correctness: jax.vmap(mixmatch_loss) matches lax.map; checks finiteness.
• Permutation invariance: shuffling labeled and unlabeled batches doesn’t change the loss; checks finiteness.
• Numerical stability (extreme logits): loss finite for huge logits; gradients w.r.t. labeled logits and unlabeled logits are finite.
• lambda_u = 0 supervised-only: returns supervised cross-entropy when unsupervised weight is zero.
• Stop-gradient on unlabeled targets: gradient w.r.t. unlabeled_targets is zero (targets are treated as constants).
• Unsupervised term zero when targets match probs: if unlabeled_targets == softmax(unlabeled_logits), unsupervised loss becomes ~0, so total ≈ supervised-only.
• bfloat16 run: smoke test that it runs in bfloat16 and returns finite output.

[surajyadav-research]

Hi @rdyro,
Could you please review this implementation when you have time? I included few tests to check robustness of losses; let me know if any of them seem unnecessary and I’ll remove them.

[vroulet]

Looks pretty good, thanks. Just add references and use correct headers

[vroulet]

2026

[vroulet]

Add reference to paper (be careful about formatting, see e.g. how it is done in the docstring of adam)

[vroulet]

Same here, add reference

[vroulet]

Use the wrapper with self.subTest(...) (see tests in other files).
This way if one test fails the other one can still be tested (so we get all info at once).

Update all tests with that pattern

[vroulet]

Make these functions not class functions but private functions at the fiel level since they are used in both tests I believe

[surajyadav-research]

@vroulet Thank you for reviewing the code. I’ll push the updated changes ASAP.

[surajyadav-research]

Hi @vroulet,
I’ve updated all the changes. Whenever you have time, could you please review them?