Don't create a mesh inside jit in sharded-computation tutorial

PiperOrigin-RevId: 707288222

Author: yashk2810

docs/sharded-computation.ipynb

@@ -399,7 +399,6 @@
     "@jax.jit\n",
     "def f_contract_2(x):\n",
     "  out = x.sum(axis=0)\n",
-    "  mesh = jax.make_mesh((8,), ('x',))\n",
     "  sharding = jax.sharding.NamedSharding(mesh, P('x'))\n",
     "  return jax.lax.with_sharding_constraint(out, sharding)\n",
     "\n",

docs/sharded-computation.md

@@ -143,7 +143,6 @@ For example, suppose that within `f_contract` above, you'd prefer the output not
 @jax.jit
 def f_contract_2(x):
   out = x.sum(axis=0)
-  mesh = jax.make_mesh((8,), ('x',))
   sharding = jax.sharding.NamedSharding(mesh, P('x'))
   return jax.lax.with_sharding_constraint(out, sharding)