Updated bridge guide

Author: vfdev-5

docs_nnx/guides/bridge_guide.ipynb

@@ -36,7 +36,7 @@
     "from flax.nnx import bridge\n",
     "import jax\n",
     "from jax import numpy as jnp\n",
-    "from jax.experimental import mesh_utils\n",
+    "from jax.sharding import PartitionSpec as P, NamedSharding, AxisType\n",
     "from typing import *"
    ]
   },
@@ -686,15 +686,14 @@
     "\n",
     "\n",
     "print(f'We have {len(jax.devices())} fake JAX devices now to partition this model...')\n",
-    "mesh = jax.sharding.Mesh(devices=mesh_utils.create_device_mesh((2, 4)),\n",
-    "                         axis_names=('in', 'out'))\n",
+    "mesh = jax.make_mesh((2, 4), ('in', 'out'), axis_types=(AxisType.Auto, AxisType.Auto))\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",
-    "print(model.w.sharding)        # The partition annotation attached with `w`\n",
-    "print(model.w.value.sharding)  # The underlying JAX array is sharded across the 2x4 mesh"
+    "print(type(model.w))                 # `nnx.Param`\n",
+    "print(model.w.sharding)              # The partition annotation attached with `w`\n",
+    "print(model.w.get_value().sharding)  # The underlying JAX array is sharded across the 2x4 mesh"
    ]
   },
   {
@@ -703,9 +702,9 @@
    "metadata": {},
    "source": [
     "    We have 8 fake JAX devices now to partition this model...\n",
-    "    <class 'flax.nnx.variables.Param'>\n",
-    "    ('in', 'out')\n",
-    "    GSPMDSharding({devices=[2,4]<=[8]})"
+    "    <class 'flax.nnx.variablelib.Param'>\n",
+    "    NamedSharding(mesh=Mesh('in': 2, 'out': 4, axis_types=(Auto, Auto)), spec=PartitionSpec('in', 'out'), memory_kind=device)\n",
+    "    NamedSharding(mesh=Mesh('in': 2, 'out': 4, axis_types=(Auto, Auto)), spec=PartitionSpec('in', 'out'), memory_kind=device)"
    ]
   },
   {
@@ -737,8 +736,9 @@
    "source": [
     "class NNXDotWithParititioning(nnx.Module):\n",
     "  def __init__(self, in_dim: int, out_dim: int, rngs: nnx.Rngs):\n",
-    "    init_fn = nnx.with_partitioning(nnx.initializers.lecun_normal(), ('in', 'out'))\n",
-    "    self.w = nnx.Param(init_fn(rngs.params(), (in_dim, out_dim)))\n",
+    "    init_fn = nnx.initializers.lecun_normal()\n",
+    "    self.w = nnx.Param(init_fn(rngs.params(), (in_dim, out_dim)),\n",
+    "                       sharding_names=('in', 'out'))\n",
     "  def __call__(self, x: jax.Array):\n",
     "    return x @ self.w\n",
     "\n",
@@ -751,7 +751,7 @@
     "  # A `NNXMeta` wrapper of the underlying `nnx.Param`\n",
     "  assert type(variables['params']['w']) == bridge.NNXMeta\n",
     "  # The annotation coming from the `nnx.Param` => (in, out)\n",
-    "  assert variables['params']['w'].metadata['sharding'] == ('in', 'out')\n",
+    "  assert variables['params']['w'].metadata['sharding_names'] == ('in', 'out')\n",
     "\n",
     "  unboxed_variables = nn.unbox(variables)\n",
     "  variable_pspecs = nn.get_partition_spec(variables)\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",
@@ -774,7 +774,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "    GSPMDSharding({devices=[2,4]<=[8]})"
+    "    NamedSharding(mesh=Mesh('in': 2, 'out': 4, axis_types=(Auto, Auto)), spec=PartitionSpec('in', 'out'), memory_kind=device)"
    ]
   },
   {

docs_nnx/guides/bridge_guide.md

@@ -36,7 +36,7 @@ from flax import linen as nn
 from flax.nnx import bridge
 import jax
 from jax import numpy as jnp
-from jax.experimental import mesh_utils
+from jax.sharding import PartitionSpec as P, NamedSharding, AxisType
 from typing import *
 ```
 
@@ -367,21 +367,20 @@ def create_sharded_nnx_module(x):
 
 
 print(f'We have {len(jax.devices())} fake JAX devices now to partition this model...')
-mesh = jax.sharding.Mesh(devices=mesh_utils.create_device_mesh((2, 4)),
-                         axis_names=('in', 'out'))
+mesh = jax.make_mesh((2, 4), ('in', 'out'), axis_types=(AxisType.Auto, AxisType.Auto))
 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`
-print(model.w.sharding)        # The partition annotation attached with `w`
-print(model.w.value.sharding)  # The underlying JAX array is sharded across the 2x4 mesh
+print(type(model.w))                 # `nnx.Param`
+print(model.w.sharding)              # The partition annotation attached with `w`
+print(model.w.get_value().sharding)  # The underlying JAX array is sharded across the 2x4 mesh
 ```
 
     We have 8 fake JAX devices now to partition this model...
-    <class 'flax.nnx.variables.Param'>
-    ('in', 'out')
-    GSPMDSharding({devices=[2,4]<=[8]})
+    <class 'flax.nnx.variablelib.Param'>
+    NamedSharding(mesh=Mesh('in': 2, 'out': 4, axis_types=(Auto, Auto)), spec=PartitionSpec('in', 'out'), memory_kind=device)
+    NamedSharding(mesh=Mesh('in': 2, 'out': 4, axis_types=(Auto, Auto)), spec=PartitionSpec('in', 'out'), memory_kind=device)
 
 +++
 
@@ -396,8 +395,9 @@ Like with any Linen metadata wrappers, you can use `linen.unbox()` to get the ra
 ```{code-cell} ipython3
 class NNXDotWithParititioning(nnx.Module):
   def __init__(self, in_dim: int, out_dim: int, rngs: nnx.Rngs):
-    init_fn = nnx.with_partitioning(nnx.initializers.lecun_normal(), ('in', 'out'))
-    self.w = nnx.Param(init_fn(rngs.params(), (in_dim, out_dim)))
+    init_fn = nnx.initializers.lecun_normal()
+    self.w = nnx.Param(init_fn(rngs.params(), (in_dim, out_dim)),
+                       sharding_names=('in', 'out'))
   def __call__(self, x: jax.Array):
     return x @ self.w
 
@@ -410,7 +410,7 @@ def create_sharded_variables(key, x):
   # A `NNXMeta` wrapper of the underlying `nnx.Param`
   assert type(variables['params']['w']) == bridge.NNXMeta
   # The annotation coming from the `nnx.Param` => (in, out)
-  assert variables['params']['w'].metadata['sharding'] == ('in', 'out')
+  assert variables['params']['w'].metadata['sharding_names'] == ('in', 'out')
 
   unboxed_variables = nn.unbox(variables)
   variable_pspecs = nn.get_partition_spec(variables)
@@ -422,14 +422,14 @@ 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
 print(variables['params']['w'].sharding)
 ```
 
-    GSPMDSharding({devices=[2,4]<=[8]})
+    NamedSharding(mesh=Mesh('in': 2, 'out': 4, axis_types=(Auto, Auto)), spec=PartitionSpec('in', 'out'), memory_kind=device)
 
 +++