code: def _validate_mocked_process_indices(devices, one_device_per_chip):
      process_to_devices = collections.defaultdict(list)
      for d in devices:
          process_to_devices[d.process_index].append(d)
  
      for local_devices in process_to_devices.values():
          if one_device_per_chip:
              # 4 devices per process
              assert len(local_devices) == 4, local_devices
          else:
              # 8 devices per process
              assert len(local_devices) == 8, local_devices
          # All devices have same z coord
          assert (
              len({d.coords[2] for d in local_devices}) == 1
          ), local_devices
          # All devices in a 2x2 subgrid
          min_coords = min(d.coords for d in local_devices)
          expected = set()
          for x, y in [(0, 0), (0, 1), (1, 0), (1, 1)]:
              expected.add(
                  (min_coords[0] + x, min_coords[1] + y, min_coords[2])
              )
          assert {
              d.coords for d in local_devices
          } == expected, local_devices
  
imports: import collections
  from collections.abc import Sequence
  import dataclasses
  import numpy
  from absl import logging
  from absl.testing import absltest
  from absl.testing import parameterized
  from jax.experimental import mesh_utils
  from jax.sharding import Mesh
  from jax._src import test_util
label_desc: Use image node_id 2 for calling a global function with example usage: _validate_mocked_process_indices(devices, one_device_per_chip) without return types