Pull Requests google/flax

Add out_sharding argument to call methods for layers with jax calls that support it pull ready

#5102 by samanklesaria was merged 2025-12-11 23:48

Fixed deprecated .value usage failing CI tests pull ready

#5097 by vfdev-5 was merged 2025-11-18 22:32

update jax minver to 0.8.1 pull ready

#5095 by cgarciae was merged 2025-11-19 20:38

Allow substring matching in `nnx.PathContains` pull ready

#5094 by thijs-vanweezel was merged 2025-12-02 19:34

Added is_causal mask argument to flax.nnx.dot_product_attention pull ready

#5093 by ibbyml was merged 2025-12-11 21:45

fix main pull ready

#5090 by cgarciae was merged 2025-11-14 22:39

Rename sharding_names to sharding_metadata pull ready

#5089 by samanklesaria was merged 2026-01-08 22:46

Support for python 3.14 pull ready

#5087 by vfdev-5 was closed 2026-01-07 23:02

add HijaxVariable [rollforward]

#5085 by copybara-service[bot] was closed 2025-11-14 19:10

Add __eq__ for variables pull ready

#5084 by samanklesaria was merged 2025-11-17 21:55

docs: Document `allow_duplicates` argument of `nnx.to_arrays`. pull ready

#5083 by dan-zheng was merged 2025-11-11 19:08

fix main pull ready

#5081 by cgarciae was merged 2025-11-10 20:41

add promote_dtype to all standard layers pull ready

#5080 by cgarciae was merged 2025-11-11 19:40

Add nnx.use_eager_sharding context manager pull ready

#5079 by samanklesaria was merged 2025-11-12 21:36

The default AxisType returned by `make_mesh` is going to change to `AxisType.Explicit` in a follow up CL. This change makes it so that the current code falls back to the old behavior.

#5078 by copybara-service[bot] was closed 2025-11-07 22:04

Added missing LayerNorm test case into TestLayersSameGraph pull ready

#5076 by vfdev-5 was merged 2025-11-10 20:39

Added kernel_metadata/bias_metadata args to nnx layers pull ready

#5074 by vfdev-5 was merged 2025-11-12 19:07

Fixed nnx.scan with carry as pytree and sow pull ready

#5073 by vfdev-5 was merged 2025-11-17 17:36

Adding set_mode_info pull ready

#5071 by chapman20j was merged 2025-11-17 17:34

support explicit sharding in eager sharding pull ready

#5070 by cgarciae was merged 2025-11-10 19:33

add nnx.dataclass pull ready

#5066 by cgarciae was merged 2025-11-12 01:23

[JAX] Replace reference to jax._src.typing with public APIs.

#5062 by copybara-service[bot] was closed 2025-10-30 15:56

Fixing Optimizer docstring - fixing #5060 pull ready

#5061 by Lucas-Fernandes-Martins was merged 2025-10-30 22:05

Pooling operation pull ready

#5057 by jorisSchaller was merged 2025-12-09 22:04

Set Mode with Tests pull ready

#5056 by chapman20j was merged 2025-10-30 19:52

Fix bound method auto-unbinding for NNX transforms pull ready

#5055 by mohsinm-dev was merged 2025-11-17 17:44

deprecate Variable.value pull ready

#5052 by cgarciae was merged 2025-11-17 20:13

filter grad state in nnx.Optimizer

#5049 by copybara-service[bot] was merged 2025-10-23 23:22

Internal change

#5048 by copybara-service[bot] was merged 2025-10-23 21:12

Fix ValueError when `nnx.jit` is used with `nnx.custom_vjp` pull ready

#5045 by samanklesaria was merged 2025-10-24 19:05