All pull requests

Pull requests

• Allow micro_grad to reduce loss auxiliary outputs

  #1888·bingdongni opened Nov 15, 2026·

  ·
  1

• Fix negative batch axes and explicit-sharding PartitionSpec handling in optax.microbatching

  #1887·copybara-service[bot] opened Nov 14, 2026Bot·

  ·

• [optax] Preserve schedule dtypes in inject_hyperparams when params is None.

  #1886·copybara-service[bot] opened Nov 13, 2026Bot·

  ·

• Make noisy_sgd key required and eta/gamma keyword-only

  #1884·wangkangmao opened Nov 13, 2026·

  ·
  1

  1

• Accept learning rate schedules in sm3

  #1883·vineethsaivs opened Nov 12, 2026·

  ·

• Add an opt-in factorized NNLS gradient path for wide matrices

  #1882·Ayushdevo opened Nov 12, 2026·

  ·
  1

• Fix zero-scale simplex projection and input derivatives

  #1880·itsrayeedscodespace opened Nov 12, 2026·

  ·
  1

  1

• [optax] Preserve low-precision schedule dtypes in inject_hyperparams when params is None.

  #1874·copybara-service[bot] opened Nov 10, 2026Bot·

  ·

• Remove legacy optax.AdaptiveGradClipState alias for optax.EmptyState.

  #1873·copybara-service[bot] opened Nov 10, 2026Bot·

  ·

• Accept Python float inputs in log_cosh

  #1871·Daiyuanranfeng opened Nov 10, 2026·

  ·

• Preserve linear OneCycle scales at shared step boundaries

  #1870·Daiyuanranfeng opened Nov 9, 2026·

  ·
  4

• Apply piecewise schedule boundaries to unsigned steps

  #1868·satiricalguru opened Nov 9, 2026·

  ·

• Test tree-math gradients, curvature and complex moment derivatives

  #1867·Ayushdevo opened Nov 8, 2026·

  ·

• Test projection input and constraint derivatives across nested trees

  #1866·Ayushdevo opened Nov 8, 2026·

  ·

• Keep reduce_on_plateau metric state in at least float32

  #1865·ely2ba opened Nov 6, 2026·

  ·
  1

• Start warmup_constant_schedule at the peak when warmup_steps=0

  #1864·ely2ba opened Nov 6, 2026·

  ·
  1

• Accept updates exactly at max_squared_norm in skip_large_updates

  #1863·ely2ba opened Nov 6, 2026·

  ·
  1

• Add analytical gradient contracts for classification and Fenchel-Young losses

  #1862·Ayushdevo opened Nov 6, 2026·

  ·

• Repair executable optimizer and assignment documentation examples

  #1861·Ayushdevo opened Nov 6, 2026·

  ·

• Avoid intermediate overflow in representable Huber losses

  #1860·Ayushdevo opened Nov 6, 2026·

  ·

• Accept schedules in scale_by_novograd and scale_by_optimistic_gradient

  #1859·mayuriphad opened Nov 6, 2026·

  ·
  1

  3

• Return scalar arrays from skip_not_finite for inactive optimizer groups

  #1855·DiadkoShmek opened Nov 5, 2026·

  ·

• Reset Lookahead at synchronized fast parameters

  #1854·DiadkoShmek opened Nov 4, 2026·

  ·

• Avoid gradient-product underflow in Rprop sign comparisons

  #1853·DiadkoShmek opened Nov 2, 2026·

  ·
  5

• Keep zero learning rates from entering MADGRAD dual averages

  #1852·DiadkoShmek opened Nov 2, 2026·

  ·
  3