Pull Requests google/flax

Rename abstract_with_sharding to as_abstract

#5388 opened 2026-04-01 16:42 by samanklesaria

add nnx.prefix

#5387 opened 2026-04-01 01:03 by copybara-service[bot]

Deprecate the keep_rngs argument in recurrent.py

#5385 opened 2026-03-31 18:50 by copybara-service[bot]

Rename vars_as to with_vars

#5384 opened 2026-03-30 15:17 by samanklesaria

Add flag for variables in nnx.merge

#5381 opened 2026-03-27 20:55 by samanklesaria

Removed bonsai from docs

#5380 opened 2026-03-27 16:16 by vfdev-5

fix Rngs.__getattr__ returning default stream for dunder names (#5375)

#5379 opened 2026-03-27 00:57 by mohsinm-dev

Fixed broken url in the sow warning

#5377 opened 2026-03-26 13:52 by vfdev-5

Consolidate "view" functions with similar naming convention.

#5376 opened 2026-03-25 21:09 by samanklesaria

[pmap] In-line definitions of `jax.device_put_sharded` and `jax.device_put_replicated`.

#5374 opened 2026-03-25 09:26 by copybara-service[bot]

Upgrade third_party/python_runtime/v3_13_unstable to 3.13.11

#5371 opened 2026-03-24 09:54 by copybara-service[bot]

Update pytree guide to use nnx.capture

#5370 opened 2026-03-23 15:34 by samanklesaria

Update the RNG guide with split and without StateAxes

#5363 opened 2026-03-20 17:06 by samanklesaria

Remove spmd_lowering from MeshComputation and do more cleanups after pmap, PmapSharding is gone

#5360 opened 2026-03-20 06:09 by copybara-service[bot]

improve scaling up guide

#5354 opened 2026-03-18 19:30 by cgarciae

Rename `SDS(vma: frozenset = ...)` to `SDS(manual_type: jax.sharding.ManualAxisType = ...)`

#5351 opened 2026-03-18 02:16 by copybara-service[bot]

Add 'prefix' argument to temp_flip_flag function

#5343 opened 2026-03-17 20:28 by samanklesaria

Fix broken error-message formatting in flax.nnx

#5334 opened 2026-03-16 13:35 by gdrosa

Add nnx.prefix

#5325 opened 2026-03-12 16:14 by samanklesaria

add transforms_tree guide

#5307 opened 2026-03-04 18:22 by cgarciae

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