code: def create_device_mesh(config: default.Config):
      """Creates a device mesh with each slice in its own data parallel group. If there is only one slice, uses two replicas."""
      devices = jax.devices()
      num_devices = len(devices)
      try:
          num_slices = 1 + max([d.slice_index for d in devices])
      except:
          num_slices = 1
      num_devices_per_slice = num_devices // num_slices
      logging.info(f"Devices: {devices}")
      logging.info(f"Number of devices: {num_devices}")
  
      multi_slice_env = hasattr(jax.devices()[0], "slice_index")
  
      dcn_parallelism = [
          config.dcn_data_parallelism,
          config.dcn_fsdp_parallelism,
          config.dcn_tensor_parallelism,
      ]
      ici_parallelism = [
          config.ici_data_parallelism,
          config.ici_fsdp_parallelism,
          config.ici_tensor_parallelism,
      ]
  
      # Find possible unspecified parallelisms
      dcn_parallelism = fill_unspecified_mesh_axes(
          dcn_parallelism, num_slices, "DCN"
      )
      ici_parallelism = fill_unspecified_mesh_axes(
          ici_parallelism, num_devices_per_slice, "ICI"
      )
  
      if multi_slice_env:
          mesh = mesh_utils.create_hybrid_device_mesh(
              ici_parallelism, dcn_parallelism
          )
      else:
          mesh = mesh_utils.create_device_mesh(ici_parallelism)
  
      logging.info(f"Decided on mesh: {mesh}")
      logging.info(f"Mesh shape: {mesh.shape}")
  
      return mesh
  
file_path: flax/flax/experimental/nnx/examples/lm1b/utils.py
imports: import logging
  from typing import Any, Callable
  from typing_extensions import Protocol, runtime_checkable
  import jax
  import jax.numpy
  import numpy
  from jax.experimental import mesh_utils
  from flax.experimental.nnx.examples.lm1b.configs import default
  from models import TransformerConfig, TransformerLM
  from flax.experimental import nnx
  from flax.training import train_state
label_desc: Use image node_id 1 for calling a global function with example usage: create_device_mesh(config) and returns: mesh