All pull requests

Pull requests

• Grow DynamicScale after growth_interval finite steps, not one more

  #5614·vineethsaivs opened Oct 6, 2026·

  ·

• Honor mask in nnx.metrics.Welford

  #5613·vineethsaivs opened Oct 6, 2026·1 task done·

  ·

• Pool masked norm statistics across devices by unmasked count

  #5611·vineethsaivs opened Oct 5, 2026·1 task done·

  ·

• Make CIRCULAR ConvTranspose with transpose_kernel the transpose of Conv

  #5610·vineethsaivs opened Oct 5, 2026·1 task done·

  ·

• Initialize the nnx MultiHeadAttention cache in set_view

  #5609·vineethsaivs opened Oct 3, 2026·1 task done·

  ·

• Init nnx fused gate kernels one gate at a time

  #5608·vineethsaivs opened Oct 2, 2026·1 task done·

  ·

• Repeat nnx.GroupNorm stats along the group axis

  #5606·vineethsaivs opened Oct 1, 2026·1 task done·

  ·

• docs: add Simplified Chinese translation (README.zh.md)

  #5605·JasonYeYuhe opened Sep 29, 2026·2 of 4 tasks·

  ·

• Sowlib scan

  #5604·samanklesaria opened Sep 28, 2026Collaborator·

  ·

• Default nnx.EMA to Variables with inexact values

  #5603·Arthur031221 opened Sep 28, 2026·3 tasks done·

  ·
  1

  4

• docs: fix 40 dead links in the Linen docs and README

  #5602·pratikgx opened Sep 28, 2026·

  ·
  1

• fix(nnx): ensure ToLinen args and kwargs are hashable

  #5598·reginaldalfret opened Sep 23, 2026·

  ·
  1

• Add the missing b_hn bias to nnx.GRUCell

  #5595·vineethsaivs opened Sep 21, 2026·

  ·
  1

  2

• fix(traverse_util): raise ValueError when flatten_dict separator conflicts with keys

  #5593·mayuriphad opened Sep 21, 2026·

  ·
  1

  2

• Fix NNX MultiHeadAttention is_causal mask during decode

  #5592·Anu27n opened Sep 20, 2026·4 tasks done·

  ·
  1

  3

• Fix NNX axis metadata under Linen scan and vmap

  #5590·Junyi-Zheng opened Sep 20, 2026·3 of 4 tasks·

  ·
  1

• Count broadcast mask entries in nnx.metrics.Average

  #5588·Nicholas022400701 opened Sep 17, 2026·2 of 4 tasks·

  ·
  2

• Allow nnx.List modules to be used with nnx.capture

  #5585·samanklesaria opened Sep 15, 2026Collaborator·

  ·
  1

  2

• Add capture_noop argument to module.sow

  #5584·samanklesaria opened Sep 14, 2026Collaborator·

  ·
  2

• Suppress pyrefly errors related to JAX scalar types

  #5583·copybara-service[bot] opened Sep 14, 2026Bot·

  ·

• fix(nnx): skip Sequence modules in capture to prevent mixed key types (fixes #5571)

  #5582·jdymitarai opened Sep 10, 2026·

  ·
  1

  8

• fix(nnx): store WeightNorm scales as Param so they train

  #5579·Alphaxiaoteng opened Sep 3, 2026·1 of 2 tasks·

  ·
  1

  3

• Add basic functional sow

  #5575·samanklesaria opened Aug 28, 2026Collaborator·

  ·

• Initial commit

  #5574·nkkongnavin-ai opened Aug 28, 2026·4 tasks done·

  ·
  1

• [No functional change] Delete VJPHiPrimitive and replace with HiPrim. This is a simple replace.

  #5567·copybara-service[bot] opened Aug 21, 2026Bot·

  ·