-
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
Loading…
2 of 4 tasks
[bridge module] Add
bridge.share_scope for layer-sublayer pairs.
#4638
opened Mar 19, 2025 by
copybara-service
bot
Loading…
[bridge module] Allow name arg to represent actual submodule path
pull ready
#4634
opened Mar 18, 2025 by
IvyZX
Loading…
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
Previous Next
ProTip!
Type g i on any issue or pull request to go back to the issue listing page.