Pull Requests google/flax

raise on Variable inplace operators pull ready

#4967 by cgarciae was merged 2025-09-24 20:06

import PathParts from flax.typing pull ready

#4966 by cgarciae was merged 2025-09-19 23:53

add Variable.set_metadata

#4965 by cgarciae was merged 2025-09-19 18:49

raise on Variable inplace operators

#4964 by cgarciae was merged 2025-09-19 18:48

fix PathParts import

#4963 by cgarciae was merged 2025-09-19 18:48

fix Variable pytree flatten

#4962 by copybara-service[bot] was merged 2025-09-19 18:03

Move iter_children and iter_modules to functions pull ready

#4961 by samanklesaria was merged 2025-10-03 18:53

Hijax Variable pull ready

#4960 by cgarciae was closed 2025-11-07 23:43

Minor typo fix in nnx.call docstring pull ready

#4959 by vfdev-5 was merged 2025-09-18 22:25

allow split tuples in Rngs.fork pull ready

#4958 by cgarciae was merged 2025-09-18 22:33

Fixed nnx.is_data docstring rendering pull ready

#4957 by vfdev-5 was merged 2025-09-18 18:08

Flax RNG guide for jax.jit: clarify rng outputs are shared but not inputs.

#4956 by copybara-service[bot] was merged 2025-09-19 12:12

Explicitly cast attribute names to string before checking for private attributes.

#4955 by copybara-service[bot] was merged 2025-09-18 21:18

improve Variable flatten pull ready

#4954 by cgarciae was merged 2025-09-18 22:21

add flax_hijax_variable flag pull ready

#4953 by cgarciae was merged 2025-09-18 21:46

fix toy examples pull ready

#4952 by cgarciae was merged 2025-09-18 18:16

expose pytree guide pull ready

#4951 by cgarciae was merged 2025-09-18 18:16

Add reference of flip doc to gspmd guide pull ready

#4949 by IvyZX was merged 2025-09-16 22:28

Add compute_flops and compute_vjp_flops options to `nnx.tabulate` pull ready

#4948 by samanklesaria was merged 2025-09-29 18:17

always differentiate standalone Variables in nnx.grad pull ready

#4947 by cgarciae was merged 2025-09-15 23:22

only copy dicts on State.__getitem__ pull ready

#4946 by cgarciae was merged 2025-09-15 21:08

add variables argument to nnx.clone pull ready

#4945 by cgarciae was merged 2025-09-15 20:56

Simplify the effort to create sharded Flax NNX models.

#4941 by copybara-service[bot] was closed 2025-10-10 01:16

Clarify `*Norm` layer docstrings: `axis_index_groups` is unused under SPMD jit.

#4940 by copybara-service[bot] was merged 2025-09-26 07:15

Implement instance norm in NNX pull ready

#4939 by mattbahr was merged 2025-09-15 23:24

update bridge wrappers from maxtext pull ready

#4937 by cgarciae was merged 2025-09-18 23:57

fix HashableMapping hash definition for mixed key types

#4936 by copybara-service[bot] was merged 2025-09-19 00:04

Fix doctest

#4935 by vfdev-5 was closed 2025-09-09 15:19

Ported nnx.PReLU from linen pull ready

#4934 by vfdev-5 was merged 2025-09-12 19:09

fix mypy

#4931 by copybara-service[bot] was merged 2025-09-09 14:48