-
Notifications
You must be signed in to change notification settings - Fork 683
Pull requests: google/flax
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
Pull requests list
Apply vmap jnp.pad over the batch in REFLECT convolution mode
#4640
opened Mar 20, 2025 by
sarlinpe
2 of 4 tasks
[bridge module] Add
bridge.share_scope for layer-sublayer pairs.
#4638
opened Mar 19, 2025 by
copybara-service
bot
Implement spectral and instance norm in NNX
#4623
opened Mar 14, 2025 by
mattbahr
Loading…
3 of 4 tasks
Stop passing reduce_axes to jax.grad, jax.vjp, and jax.value_and_grad.
#4617
opened Mar 11, 2025 by
copybara-service
bot
Loading…
Adding nnx Gemma2-2b (including overall fixes) to examples/gemma
#4587
opened Feb 28, 2025 by
mdda
Loading…
Canonicalize PartitionSpec so that we can delete ParsedPartitionSpec
#4577
opened Feb 26, 2025 by
copybara-service
bot
Loading…
Add some methods on Sequential container
#4565
opened Feb 22, 2025 by
Ruhaan838
Loading…
2 of 3 tasks
Fixes #1925 non-str dict keys not suppoted in module state
#4563
opened Feb 21, 2025 by
muhrin
Loading…
3 of 4 tasks
Improve error message for when nnx.Modules use jax or numpy arrays as leaf values
#4492
opened Jan 20, 2025 by
RaghuSpaceRajan
Loading…
[linen] Linesearch (and lbfgs) support for TrainState
#4471
opened Jan 6, 2025 by
emiresenov
Loading…
3 of 4 tasks
Update NNX
State.state.merge method docs in statelib.py
#4453
opened Dec 17, 2024 by
8bitmp3
Loading…
Update NNX
State.state.filter method docs in statelib.py
#4452
opened Dec 17, 2024 by
8bitmp3
Loading…
Previous Next
ProTip!
Updated in the last three days: updated:>2025-03-18.