Merge pull request #5149 from IvyZX:release

Flax Authors · Flax Authors · commit d1791844e97b · 2025-12-18T14:05:40.000-08:00
PiperOrigin-RevId: 846407227
diff --git a/README.md b/README.md
@@ -158,7 +158,7 @@ To cite this repository:
   author = {Jonathan Heek and Anselm Levskaya and Avital Oliver and Marvin Ritter and Bertrand Rondepierre and Andreas Steiner and Marc van {Z}ee},
   title = {{F}lax: A neural network library and ecosystem for {JAX}},
   url = {http://github.com/google/flax},
-  version = {0.12.1},
+  version = {0.12.2},
   year = {2024},
 }
 ```
diff --git a/docs/guides/parallel_training/flax_on_pjit.ipynb b/docs/guides/parallel_training/flax_on_pjit.ipynb
@@ -669,7 +669,7 @@
     "  state = state.apply_gradients(grads=grads)\n",
     "  return state\n",
     "\n",
-    "with mesh:\n",
+    "with jax.set_mesh(mesh):\n",
     "  new_state = train_step(initialized_state, x)"
    ]
   },
@@ -824,7 +824,7 @@
     "def apply_fn(state, x):\n",
     "  return state.apply_fn({'params': state.params}, x)\n",
     "\n",
-    "with mesh:\n",
+    "with jax.set_mesh(mesh):\n",
     "  y = apply_fn(new_state, x)\n",
     "print(type(y))\n",
     "print(y.dtype)\n",
@@ -861,7 +861,7 @@
     "  jax.tree_util.tree_map(lambda x: x.block_until_ready(), xs)\n",
     "  return xs\n",
     "\n",
-    "with mesh:\n",
+    "with jax.set_mesh(mesh):\n",
     "  new_state = block_all(train_step(initialized_state, x))"
    ]
   },
diff --git a/docs/guides/parallel_training/flax_on_pjit.md b/docs/guides/parallel_training/flax_on_pjit.md
@@ -319,7 +319,7 @@ def train_step(state, x):
   state = state.apply_gradients(grads=grads)
   return state
 
-with mesh:
+with jax.set_mesh(mesh):
   new_state = train_step(initialized_state, x)
 ```
 
@@ -338,7 +338,7 @@ Then, create a compiled inference step. Note that the output is also sharded alo
 def apply_fn(state, x):
   return state.apply_fn({'params': state.params}, x)
 
-with mesh:
+with jax.set_mesh(mesh):
   y = apply_fn(new_state, x)
 print(type(y))
 print(y.dtype)
@@ -357,7 +357,7 @@ def block_all(xs):
   jax.tree_util.tree_map(lambda x: x.block_until_ready(), xs)
   return xs
 
-with mesh:
+with jax.set_mesh(mesh):
   new_state = block_all(train_step(initialized_state, x))
 ```
 
diff --git a/docs_nnx/flip/4844-var-eager-sharding.md b/docs_nnx/flip/4844-var-eager-sharding.md
@@ -33,7 +33,7 @@ def create_sharded_model():
   return model
 
 mesh = jax.make_mesh(((2, 4)), ("data", "model"))
-with mesh:
+with jax.set_mesh(mesh):
   sharded_model = create_sharded_model()
 ```
 
diff --git a/docs_nnx/guides/bridge_guide.ipynb b/docs_nnx/guides/bridge_guide.ipynb
@@ -689,7 +689,7 @@
     "mesh = jax.sharding.Mesh(devices=mesh_utils.create_device_mesh((2, 4)),\n",
     "                         axis_names=('in', 'out'))\n",
     "x = jax.random.normal(jax.random.key(42), (4, 32))\n",
-    "with mesh:\n",
+    "with jax.set_mesh(mesh):\n",
     "  model = create_sharded_nnx_module(x)\n",
     "\n",
     "print(type(model.w))           # `nnx.Param`\n",
@@ -763,7 +763,7 @@
     "                              nn.get_partition_spec(variables))\n",
     "  return sharded_vars\n",
     "\n",
-    "with mesh:\n",
+    "with jax.set_mesh(mesh):\n",
     "  variables = create_sharded_variables(jax.random.key(0), x)\n",
     "\n",
     "# The underlying JAX array is sharded across the 2x4 mesh\n",
diff --git a/docs_nnx/guides/bridge_guide.md b/docs_nnx/guides/bridge_guide.md
@@ -370,7 +370,7 @@ print(f'We have {len(jax.devices())} fake JAX devices now to partition this mode
 mesh = jax.sharding.Mesh(devices=mesh_utils.create_device_mesh((2, 4)),
                          axis_names=('in', 'out'))
 x = jax.random.normal(jax.random.key(42), (4, 32))
-with mesh:
+with jax.set_mesh(mesh):
   model = create_sharded_nnx_module(x)
 
 print(type(model.w))           # `nnx.Param`
@@ -422,7 +422,7 @@ def create_sharded_variables(key, x):
                               nn.get_partition_spec(variables))
   return sharded_vars
 
-with mesh:
+with jax.set_mesh(mesh):
   variables = create_sharded_variables(jax.random.key(0), x)
 
 # The underlying JAX array is sharded across the 2x4 mesh
diff --git a/examples/lm1b/utils.py b/examples/lm1b/utils.py
@@ -168,7 +168,7 @@ def setup_initial_state(model, tx, config, rng, mesh):
   state_logical_annotations = nn.get_partition_spec(abstract_state)
 
   # Initialization
-  with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+  with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):
     state_mesh_annotations = nn.logical_to_mesh_sharding(
         state_logical_annotations, mesh, config.logical_axis_rules
     )
diff --git a/flax/core/meta.py b/flax/core/meta.py
@@ -30,7 +30,6 @@
 from flax import errors, struct
 from flax.typing import LogicalNames
 import jax
-from jax.interpreters import pxla
 
 A = TypeVar('A')
 B = TypeVar('B')
@@ -182,9 +181,7 @@ def inner_update(c, v):
 def get_global_mesh() -> jax.sharding.AbstractMesh | jax.sharding.Mesh | None:
   mesh = jax.sharding.get_abstract_mesh()
   if mesh.empty:
-    mesh = pxla.thread_resources.env.physical_mesh
-    if mesh.empty:
-      return None
+    return None
   return mesh
 
 
diff --git a/flax/version.py b/flax/version.py
@@ -13,4 +13,4 @@
 # limitations under the License.
 
 """Current Flax version at head on Github."""
-__version__ = '0.12.1'
+__version__ = '0.12.2'
diff --git a/tests/nnx/module_test.py b/tests/nnx/module_test.py
@@ -693,6 +693,9 @@ def __init__(self, din, dout, *, rngs: nnx.Rngs):
       nnx.set_mode(block, deterministic=True, use_running_average=True, unknown=True)
 
   def test_cloud_pickle(self):
+    import platform
+    if platform.python_version().startswith('3.11'):
+      self.skipTest("Cloudpickle cannot pickle PRNGKeyArray on python 3.11")
     class Model(nnx.Module):
       def __init__(self, din, dmid, dout, rngs: nnx.Rngs):
         self.linear = nnx.Linear(din, dmid, rngs=rngs)

Original file line number	Diff line number	Diff line change
`@@ -158,7 +158,7 @@ To cite this repository:`
`158`	`158`	`author = {Jonathan Heek and Anselm Levskaya and Avital Oliver and Marvin Ritter and Bertrand Rondepierre and Andreas Steiner and Marc van {Z}ee},`
`159`	`159`	`title = {{F}lax: A neural network library and ecosystem for {JAX}},`
`160`	`160`	`url = {http://github.com/google/flax},`
`161`		`- version = {0.12.1},`
	`161`	`+ version = {0.12.2},`
`162`	`162`	`year = {2024},`
`163`	`163`	`}`
`164`	`164`	```
Original file line number	Diff line number	Diff line change
`@@ -168,7 +168,7 @@ def setup_initial_state(model, tx, config, rng, mesh):`
`168`	`168`	`state_logical_annotations = nn.get_partition_spec(abstract_state)`
`169`	`169`
`170`	`170`	`# Initialization`
`171`		`- with mesh, nn_partitioning.axis_rules(config.logical_axis_rules):`
	`171`	`+ with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):`
`172`	`172`	`state_mesh_annotations = nn.logical_to_mesh_sharding(`
`173`	`173`	`state_logical_annotations, mesh, config.logical_axis_rules`
`174`	`174`	`)`