All pull requests

Pull requests

• Take the absolute difference in triplet_margin_loss

  #1794·vineethsaivs opened Sep 24, 2026·

  ·

• Stabilize integer-label softmax cross-entropy gradients under high model confidence.

  #1793·copybara-service[bot] opened Sep 24, 2026Bot·

  ·

• Stabilize integer-label cross entropy under logit shifts

  #1792·sylvesterkaczmarek opened Sep 22, 2026·

  ·

• Stop adafactor's step_offset from producing a non-finite decay rate

  #1791·vineethsaivs opened Sep 21, 2026·

  ·

• Fix float16 zero-gradient NaNs in AdaDelta, Adan, and NovoGrad

  #1790·Anu27n opened Sep 18, 2026·

  ·
  1

  1

• Fix sam with sync_period=1 undoing the previous step in transparent mode

  #1788·Nicholas022400701 opened Sep 18, 2026·

  ·
  2

• Fix the weight decay update in momo and momo_adam to match the paper and the reference implementation

  #1787·Nicholas022400701 opened Sep 18, 2026·

  ·

• Keep kl_divergence finite when target probabilities are zero

  #1786·Nicholas022400701 opened Sep 17, 2026·

  ·

• Fix float16 zero-gradient NaNs in RMS transforms

  #1785·vishesh9131 opened Sep 15, 2026·

  ·
  1

• Ensure CTC loss is nonnegative against float32 rounding errors (fixes #1771)

  #1783·jdymitarai opened Sep 15, 2026·

  ·

• Fix onecycle and warmup cosine schedule docstrings to match the code

  #1781·rishabhsai opened Sep 12, 2026·

  ·

• Fix AdamW example in getting started notebook

  #1780·rishabhsai opened Sep 11, 2026·

  ·
  4

  

• Fix Pyrefly missing-attribute and unsupported-operation ignores on ArrayLike arguments across loss functions, transforms, linesearches, and schedules by normalizing them to jax.Array using jnp.asarray.

  #1776·copybara-service[bot] opened Sep 11, 2026Bot·

  ·
  1

• Preserve plateau scale dtype with array minimum scales

  #1774·sylvesterkaczmarek opened Sep 10, 2026·

  ·

• Fix Pyrefly missing-attribute and unsupported-operation ignores on ArrayLike values by normalizing them to jax.Array using jnp.asarray.

  #1773·copybara-service[bot] opened Sep 8, 2026Bot·

  ·
  1

• Fix negative CTC loss for confident alignments in float32

  #1772·rajasekharporeddy opened Sep 7, 2026Collaborator·

  ·
  1

  1

• fix(transform): prevent NaN in Adam-family optimizers on float16 zero-gradient steps

  #1769·somuai opened Sep 3, 2026·

  ·
  1

  5

• fix(transform): prevent NaN in scale_by_adamax on float16 zero-gradient steps

  #1768·somuai opened Sep 3, 2026·

  ·
  3

• Fix contrib.reduce_on_plateau never reducing the learning rate on negative values

  #1767·Kayvan-Zahiri opened Sep 2, 2026·

  ·
  1

• Fix a NaN in contrib.dadapt_adamw on zero denominators and zero learning rates.

  #1766·copybara-service[bot] opened Aug 30, 2026Bot·

  ·

• Guard contrib.prodigy against zero denominators and zero learning rates.

  #1765·copybara-service[bot] opened Aug 30, 2026Bot·

  ·

• fix(losses): avoid NaN gradient in weighted_logsoftmax at -inf logits

  #1762·jaideeppyne opened Aug 25, 2026·

  ·

• Fix all-NaN output for complex inputs with 'schatten' preconditioning in Muon

  #1761·shoemoney opened Aug 24, 2026·

  ·
  5

• Fix masking in label smoothing

  #1760·tandede opened Aug 20, 2026·

  ·
  2

• docs: fix typos and update docstring formatting across multiple modules

  #1759·rajasekharporeddy opened Aug 19, 2026Collaborator·

  ·
  1