Merge pull request #28027 from jax-ml:explicit_axes

PiperOrigin-RevId: 747889283

Author: Google-ML-Automation

docs/notebooks/explicit-sharding.ipynb

@@ -59,7 +59,7 @@
     "import numpy as np\n",
     "import jax.numpy as jnp\n",
     "from jax.sharding import PartitionSpec as P, AxisType, set_mesh, get_abstract_mesh\n",
-    "from jax.experimental.shard import reshard, auto_axes\n",
+    "from jax.experimental.shard import reshard, auto_axes, explicit_axes\n",
     "\n",
     "jax.config.update('jax_num_cpu_devices', 8)"
    ]
@@ -652,7 +652,51 @@
     "id": "_3sfJjRq8w9f"
    },
    "source": [
-    "As you can see, inside `g`, the type of `arr1` is `ShapedArray(float32[4,4@Y])` which indicates it's Explicit over `Y` mesh axis while auto over `X`."
+    "As you can see, inside `g`, the type of `arr1` is `ShapedArray(float32[4,4@Y])` which indicates it's Explicit over `Y` mesh axis while auto over `X`.\n",
+    "\n",
+    "\n",
+    "You can also use the `explicit_axes` API to drop into `Explicit` mode over some or all mesh axes."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a102e9c7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "auto_mesh = jax.make_mesh((2, 4), (\"X\", \"Y\"),\n",
+    "                           axis_types=(AxisType.Auto, AxisType.Auto))\n",
+    "\n",
+    "@functools.partial(explicit_axes, axes=('X', 'Y'))\n",
+    "def explicit_g(y):\n",
+    "  print(f'mesh inside g: {get_abstract_mesh()}')\n",
+    "  print(f'y.sharding inside g: {jax.typeof(y) = }')\n",
+    "  z = y * 2\n",
+    "  print(f'z.sharding inside g: {jax.typeof(z) = }', end='\\n\\n')\n",
+    "  return z\n",
+    "\n",
+    "@jax.jit\n",
+    "def f(arr1):\n",
+    "  print(f'mesh inside f: {get_abstract_mesh()}', end='\\n\\n')\n",
+    "  x = jnp.sin(arr1)\n",
+    "\n",
+    "  z = explicit_g(x, in_shardings=P(\"X\", \"Y\"))\n",
+    "\n",
+    "  return z + 1\n",
+    "\n",
+    "with jax.sharding.use_mesh(auto_mesh):\n",
+    "  some_x = jax.device_put(np.arange(16).reshape(4, 4), P(\"X\", \"Y\"))\n",
+    "  f(some_x)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e64d40de",
+   "metadata": {},
+   "source": [
+    "As you can see, all axes of mesh inside `f` are of type `Auto` while inside `g`, they are of type `Explicit`.\n",
+    "Because of that, sharding is visible on the type of arrays inside `g`."
    ]
   },
   {

docs/notebooks/explicit-sharding.md

@@ -56,7 +56,7 @@ import jax
 import numpy as np
 import jax.numpy as jnp
 from jax.sharding import PartitionSpec as P, AxisType, set_mesh, get_abstract_mesh
-from jax.experimental.shard import reshard, auto_axes
+from jax.experimental.shard import reshard, auto_axes, explicit_axes
 
 jax.config.update('jax_num_cpu_devices', 8)
 ```
@@ -403,6 +403,38 @@ f(some_x)
 
 As you can see, inside `g`, the type of `arr1` is `ShapedArray(float32[4,4@Y])` which indicates it's Explicit over `Y` mesh axis while auto over `X`.
 
+
+You can also use the `explicit_axes` API to drop into `Explicit` mode over some or all mesh axes.
+
+```{code-cell} ipython3
+auto_mesh = jax.make_mesh((2, 4), ("X", "Y"),
+                           axis_types=(AxisType.Auto, AxisType.Auto))
+
+@functools.partial(explicit_axes, axes=('X', 'Y'))
+def explicit_g(y):
+  print(f'mesh inside g: {get_abstract_mesh()}')
+  print(f'y.sharding inside g: {jax.typeof(y) = }')
+  z = y * 2
+  print(f'z.sharding inside g: {jax.typeof(z) = }', end='\n\n')
+  return z
+
+@jax.jit
+def f(arr1):
+  print(f'mesh inside f: {get_abstract_mesh()}', end='\n\n')
+  x = jnp.sin(arr1)
+
+  z = explicit_g(x, in_shardings=P("X", "Y"))
+
+  return z + 1
+
+with jax.sharding.use_mesh(auto_mesh):
+  some_x = jax.device_put(np.arange(16).reshape(4, 4), P("X", "Y"))
+  f(some_x)
+```
+
+As you can see, all axes of mesh inside `f` are of type `Auto` while inside `g`, they are of type `Explicit`.
+Because of that, sharding is visible on the type of arrays inside `g`.
+
 +++ {"id": "sJcWbfAh7UcO"}
 
 ## Concrete array shardings can mention `Auto` mesh axis