All pull requests

Pull requests

• Avoid recomputing jax.nn.softmax(x, axis=-1) from scratch in optax.losses._classification._weighted_logsoftmax_jvp by reusing jnp.exp(logsoftmax_x) from the primal log_softmax pass (1.52x speedup on TPU v5e: 702.0 us -> 460.9 us).

  #1849·copybara-service[bot] opened Oct 30, 2026Bot·

  ·
  1

• Compute ntxent pairwise cosine similarities via unit normalization and 2D jnp.matmul (O(N^2 D) on the MXU) instead of 3D [N, N, D] broadcast cosine_similarity + [N, 2N] concatenation, and formulate the per-pair loss with jax.nn.logsumexp and jax.nn.softplus (3.10x speedup on TPU v5e: 953.8 us -> 307.5 us, -24% VMEM).

  #1848·copybara-service[bot] opened Oct 30, 2026Bot·

  ·
  1

• Keep scalar running sums in at least float32 in schedule_free, dog, and dowg

  #1847·EvanZhuang opened Oct 29, 2026·

  ·

• Avoid intermediate overflow in L2 loss

  #1846·Linxiushen opened Oct 29, 2026·

  ·

• Preserve multiclass hinge margins under large common score offsets

  #1845·Linxiushen opened Oct 29, 2026·

  ·

• Preserve small residual values in log_cosh

  #1844·Linxiushen opened Oct 28, 2026·

  ·

• Avoid intermediate overflow in COCOB betting fraction

  #1843·Linxiushen opened Oct 28, 2026·

  ·

• Preserve optimizer state dtypes when resetting Lookahead

  #1842·Linxiushen opened Oct 28, 2026·

  ·

• Keep CTC losses and gradients finite for valid masked alignments

  #1841·salma2vec opened Oct 28, 2026·

  ·
  1

• Handle empty parameter groups in flatten

  #1839·AryanRG opened Oct 26, 2026·

  ·
  1

• Keep integer dtypes in tree_max/tree_min and return 0 for an empty inf-norm

  #1838·Sonica-B opened Oct 25, 2026·

  ·
  1

• Fix MultiSteps.has_updated after rejected mini-steps

  #1837·Linxiushen opened Oct 25, 2026·

  ·

• Fix zero-patience scheduling in reduce_on_plateau

  #1836·laasyapriya2007 opened Oct 23, 2026·

  ·
  1

• Reset non-finite MultiSteps accumulators after emission

  #1834·0-Actual opened Oct 22, 2026·

  ·
  1

• Stabilize simplex projection under large common offsets

  #1833·Linxiushen opened Oct 21, 2026·

  ·

• Fix Hermitian coefficient in complex vector projection

  #1832·Linxiushen opened Oct 21, 2026·

  ·

• Stabilize float16 RMS values and parameter block scaling

  #1831·Ayushdevo opened Oct 20, 2026·

  ·

• Compute low-precision cosine normalization in float32

  #1830·jabrailkhalil opened Oct 20, 2026·

  ·
  1

• Stabilize low-precision adaptive gradient clipping

  #1829·Linxiushen opened Oct 19, 2026·

  ·

• Preserve mixed parameter dtypes in Sophia Hessian probes

  #1828·mottopanikeiku opened Oct 19, 2026·

  ·

• Fix dice_loss rejecting valid single-channel binary probabilities

  #1826·udsy19 opened Oct 18, 2026·

  ·
  1

• Compute RAdam's sqrt denominator in float32

  #1825·rootkiller6788 opened Oct 18, 2026·

  ·

• Return zero MultiSteps updates before emission

  #1824·edawite opened Oct 18, 2026Contributor·

  ·

• docs: add tutorial for custom gradient transformations

  #1823·itbert opened Oct 14, 2026·

  ·
  1

  3

• fix(clipping): prevent float16 overflow in tree_norm and clip_by_global_norm (#1821)

  #1822·reginaldalfret opened Oct 14, 2026·

  ·
  1

  2