Go
Home
Pricing
FAQ
Install
Home
Pricing
FAQ
Install
Login
via GitHub
google/flax
Pull Requests
Commits
Open
Closed
add with_attributes
#5308 opened 2026-03-04 19:09 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]
Annotate `wrapped_fn` in `decorator_lift_transform` for trace-ability
#5298 opened 2026-03-03 03:50 by
copybara-service[bot]
add nnx_basics_tree.ipynb
#5293 opened 2026-03-02 19:16 by
cgarciae
Remove optax pin
#5292 opened 2026-03-02 15:55 by
samanklesaria
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
feat(nnx): add GQA support to MultiHeadAttention
#5259 opened 2026-02-18 15:51 by
ayulockedin
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
Add sharding propagation support in nnx.eval_shape (clone of #5111)
#5247 opened 2026-02-13 18:25 by
samanklesaria
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
Fix PyTreeNode + Generic losing __parameters__ when Generic is last in bases
#5237 opened 2026-02-10 19:29 by
mohsinm-dev
Add 'opt_sharding' option for Variables
#5231 opened 2026-02-09 20:47 by
samanklesaria
Older