Repository navigation

All pull requests

Labels (18)Milestones (0)New pull request

Open

(302)

Closed

(3,246)

Fix TrainState.apply_gradients when params is a single array
#5620aalquwayfili opened Oct 11, 20262 of 4 tasks
1
Keep masked nan and inf out of nnx.metrics.Average
#5619vineethsaivs opened Oct 9, 2026
Raise an error for unknown keys in MultiMetric.update mask dict
#5618irhyl opened Oct 9, 2026
Raise an error for MultiMetric names that shadow its attributes
#5617irhyl opened Oct 9, 2026
Grow DynamicScale after growth_interval finite steps, not one more
#5614vineethsaivs opened Oct 6, 2026
Honor mask in nnx.metrics.Welford
#5613vineethsaivs opened Oct 6, 20261 task done
Pool masked norm statistics across devices by unmasked count
#5611vineethsaivs opened Oct 5, 20261 task done
Make CIRCULAR ConvTranspose with transpose_kernel the transpose of Conv
#5610vineethsaivs opened Oct 5, 20261 task done
Initialize the nnx MultiHeadAttention cache in set_view
#5609vineethsaivs opened Oct 3, 20261 task done
Init nnx fused gate kernels one gate at a time
#5608vineethsaivs opened Oct 2, 20261 task done
Repeat nnx.GroupNorm stats along the group axis
#5606vineethsaivs opened Oct 1, 20261 task done
docs: add Simplified Chinese translation (README.zh.md)
#5605JasonYeYuhe opened Sep 29, 20262 of 4 tasks
Sowlib scan
#5604samanklesaria opened Sep 28, 2026Collaborator
Default nnx.EMA to Variables with inexact values
#5603Arthur031221 opened Sep 28, 20263 tasks done
6
docs: fix 40 dead links in the Linen docs and README
#5602pratikgx opened Sep 28, 2026
1
fix(nnx): ensure ToLinen args and kwargs are hashable
#5598reginaldalfret opened Sep 23, 2026
Add the missing b_hn bias to nnx.GRUCell
#5595vineethsaivs opened Sep 21, 2026
2
fix(traverse_util): raise ValueError when flatten_dict separator conflicts with keys
#5593mayuriphad opened Sep 21, 2026
3
Fix NNX MultiHeadAttention is_causal mask during decode
#5592Anu27n opened Sep 20, 20264 tasks done
3
Fix NNX axis metadata under Linen scan and vmap
#5590Junyi-Zheng opened Sep 20, 20263 of 4 tasks
Count broadcast mask entries in nnx.metrics.Average
#5588Nicholas022400701 opened Sep 17, 20262 of 4 tasks
2
Allow nnx.List modules to be used with nnx.capture
#5585samanklesaria opened Sep 15, 2026Collaborator
2
Add capture_noop argument to module.sow
#5584samanklesaria opened Sep 14, 2026Collaborator
2
Suppress pyrefly errors related to JAX scalar types
#5583copybara-service[bot] opened Sep 14, 2026Bot
fix(nnx): skip Sequence modules in capture to prevent mixed key types (fixes #5571)
#5582jdymitarai opened Sep 10, 2026
8

Pull requests