[JAX] Replace users of deprecated jax.lib.xla_bridge functions with m…

…ore public names.

Change in preparation for removing some function exports from jax.lib.xla_bridge.

PiperOrigin-RevId: 404849899

cloud_tpu_colabs/Pmap_Cookbook.ipynb

@@ -506,7 +506,7 @@
       },
       "source": [
         "from jax._src.lib import xla_bridge\n",
-        "device_count = xla_bridge.device_count()\n",
+        "device_count = jax.device_count()\n",
         "\n",
         "def send_right(x, axis_name):\n",
         "  left_perm = [(i, (i + 1) % device_count) for i in range(device_count)]\n",