Go
Home
Pricing
FAQ
Install
Home
Pricing
FAQ
Install
Login
via GitHub
google/flax
Pull Requests
Commits
Open
Closed
allow in/out_axes when graph_updates=False
#5323 opened 2026-03-10 01:55 by
copybara-service[bot]
fix eval_shape's _to_variable
#5321 opened 2026-03-09 17:41 by
copybara-service[bot]
add graph_updates to vmap
#5320 opened 2026-03-09 04:30 by
copybara-service[bot]
add graph_updates argument in shard_map
#5319 opened 2026-03-09 03:08 by
copybara-service[bot]
add graph_updates argument to jit
#5317 opened 2026-03-06 18:40 by
copybara-service[bot]
add tree-mode-nnx FLIP
#5310 opened 2026-03-04 22:34 by
copybara-service[bot]
add transforms_tree guide
#5307 opened 2026-03-04 18:22 by
cgarciae
Iterate over modules only through iter_children().
#5302 opened 2026-03-04 00:34 by
copybara-service[bot]
Bump the pip group across 9 directories with 4 updates
dependencies
python
#5300 opened 2026-03-04 00:07 by
dependabot[bot]
change defaults for Dropout and BatchNorm
#5299 opened 2026-03-03 14:49 by
copybara-service[bot]
add nnx_basics_tree.ipynb
#5293 opened 2026-03-02 19:16 by
cgarciae
Expand Einsum test coverage with parameterized tests
#5285 opened 2026-02-26 21:35 by
us
Add tests, docstrings and validation to metrics.py
#5284 opened 2026-02-26 21:33 by
us
Add examples and fix docstrings in nnx.nn.dtypes
#5283 opened 2026-02-26 21:11 by
us
Expand Dropout test coverage and migrate to absltest
#5281 opened 2026-02-26 20:32 by
us
Expand LoRA test coverage with dtype, gradient, and edge case tests
#5280 opened 2026-02-26 20:25 by
us
Expand TestLinenConsistency parameterized.product in attention_test
#5279 opened 2026-02-26 20:16 by
us
Expand Embed test coverage with functional, attend, and dtype tests
#5278 opened 2026-02-26 06:26 by
us
Add tests for nnx.initializers zeros_init and ones_init
#5276 opened 2026-02-25 23:17 by
us
Add tests for flax.nnx.nn.dtypes
#5275 opened 2026-02-25 22:46 by
us
Add PReLU test coverage and move TestPReLUConsistency to activations_test
#5274 opened 2026-02-25 22:07 by
us
Fixed two typos in the Module & Pytree guide
#5271 opened 2026-02-23 23:58 by
Zirunis
Add split method to RngStream
#5270 opened 2026-02-23 15:53 by
samanklesaria
Add intermediate value captures (extends #4925)
#5257 opened 2026-02-16 19:46 by
samanklesaria
feat(nnx): add out_sharding to recurrent cells
#5255 opened 2026-02-15 14:45 by
aarushisingh04
feat(nnx): add preferred_element_type to recurrent cells
#5254 opened 2026-02-15 14:45 by
aarushisingh04
feat(nnx): add preferred_element_type to attention module
#5253 opened 2026-02-15 14:45 by
aarushisingh04
Generalize out_sharding to work with NamedSharding and Format
#5246 opened 2026-02-13 02:59 by
samanklesaria
Track explicit sharding in Variables
#5245 opened 2026-02-12 23:02 by
samanklesaria
Add 'opt_sharding' option for Variables
#5231 opened 2026-02-09 20:47 by
samanklesaria
Older