Repository navigation

All pull requests

Labels (18)Milestones (0)New pull request

Open

(306)

Closed

(3,259)

Add warmup_steps support to nnx.EMA.
#5644copybara-service[bot] opened Nov 5, 2026Bot
1
[nnx] Speed up is_array_ref.
#5643copybara-service[bot] opened Nov 5, 2026Bot
1
fix(nnx): SpectralNorm/WeightNorm no longer persist normalized weights into params
#5642davrens101 opened Nov 5, 2026
1
Fix pre-existing Pyrefly type-check errors in Flax and remove migration suppressions
#5641copybara-service[bot] opened Nov 4, 2026Bot
Fix negative argnums state filters in NNX graph gradients
#5637DiadkoShmek opened Nov 2, 20262 of 4 tasks
3
ci(release): migrate PyPI publishing to OSS Exit Gate
#5636alonsopf89 opened Nov 2, 2026
Pass a PRNG key to NNX WeightNorm scale initializers
#5635DiadkoShmek opened Nov 2, 20262 tasks done
[nnx] Add BacktrackingLinesearch optimizer for line-search (#4545)
#5634XaRaaaa opened Oct 30, 20263 of 4 tasks
Preserve byte order when serializing non-native NumPy arrays
#5633AryanRG opened Oct 27, 20262 of 4 tasks
Validate feature_group_count in nnx.Conv constructor
#5624Sonica-B opened Oct 25, 20263 of 4 tasks
1
Fix mypy and doctest failures with jax 0.12.0.
#5622copybara-service[bot] opened Oct 20, 2026Bot
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
1
Repeat nnx.GroupNorm stats along the group axis
#5606vineethsaivs opened Oct 1, 20261 task done
1
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
10

Navigation Menu

Repository navigation

All pull requests

Add warmup_steps support to nnx.EMA.

[nnx] Speed up `is_array_ref`.

fix(nnx): SpectralNorm/WeightNorm no longer persist normalized weights into params

Fix pre-existing Pyrefly type-check errors in Flax and remove migration suppressions

Fix negative argnums state filters in NNX graph gradients

ci(release): migrate PyPI publishing to OSS Exit Gate

Pass a PRNG key to NNX WeightNorm scale initializers

[nnx] Add BacktrackingLinesearch optimizer for line-search (#4545)

Preserve byte order when serializing non-native NumPy arrays

Validate feature_group_count in nnx.Conv constructor

Fix mypy and doctest failures with jax 0.12.0.

Fix TrainState.apply_gradients when params is a single array

Keep masked nan and inf out of nnx.metrics.Average

Raise an error for unknown keys in MultiMetric.update mask dict

Raise an error for MultiMetric names that shadow its attributes

Grow DynamicScale after growth_interval finite steps, not one more

Honor mask in nnx.metrics.Welford

Pool masked norm statistics across devices by unmasked count

Make CIRCULAR ConvTranspose with transpose_kernel the transpose of Conv

Initialize the nnx MultiHeadAttention cache in set_view

Init nnx fused gate kernels one gate at a time

Repeat nnx.GroupNorm stats along the group axis

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

Sowlib scan

Default nnx.EMA to Variables with inexact values

Pull requests