code: def _addressable_devices_indices_map(
      sharding: Sharding, global_shape: Shape
  ) -> Mapping[Device, Index | None]:
      global_map = sharding.devices_indices_map(global_shape)
      if sharding.is_fully_addressable:
          return global_map
      if hasattr(sharding, "_internal_device_list"):
          return {
              d: global_map[d]
              for d in sharding._internal_device_list.addressable_device_list
          }
      return {
          d: ind
          for d, ind in global_map.items()
          if d.process_index == d.client.process_index()
      }
  
file_path: jax/jax/_src/sharding.py
imports: from __future__ import annotations
  from collections.abc import Mapping, Sequence
  import functools
  from jax._src import util
  from jax._src import xla_bridge
  from jax._src.lib import xla_client
label_desc: Use image node_id 1 for calling a global function with example usage: _addressable_devices_indices_map(sharding, global_shape) and returns: unknown, global_map, unknown