# Path: vmoe/partitioning.py
# def process_has_contiguous_device_slice(devices: np.ndarray,
#                                         process_index: int) -> bool:
# def processes_have_contiguous_device_slices(devices: np.ndarray) -> bool:
# def get_auto_logical_mesh(
#     num_partitions: int,
#     devices: Optional[Sequence[Device]] = None) -> Mesh:
# def get_auto_logical_mesh_other(num_partitions: int,
#                                 hardware_mesh: np.ndarray) -> Mesh:
# def get_auto_logical_mesh_tpu(num_partitions: int,
#                               hardware_mesh: np.ndarray) -> Mesh:
# def get_device_coords(device: Device) -> Union[TpuCoords, OtherCoords]:
# def get_device_coords_other(device: Device) -> OtherCoords:
# def get_device_coords_tpu(device: Device) -> TpuCoords:
# def get_hardware_mesh_local_shape(
#     local_devices: Optional[Sequence[Device]] = None) -> Tuple[int, ...]:
# def get_hardware_mesh_other(devices: Sequence[Device]) -> np.ndarray:
# def get_hardware_mesh_tpu(devices: Sequence[Device]) -> np.ndarray:
# def get_logical_mesh(partitions: Tuple[int, ...],
#                      hardware_mesh: np.ndarray) -> Mesh:
# def get_logical_mesh_default(
#     partitions: Tuple[int, ...],
#     replicas: Tuple[int, ...],
#     hardware_mesh: np.ndarray) -> Mesh:
# def get_logical_mesh_tile_by_process(
#     partitions: Tuple[int, ...],
#     replicas: Tuple[int, ...],
#     hardware_mesh: np.ndarray,
#     hardware_mesh_local_shape: Tuple[int, ...]) -> Mesh:
#   def fn(p, r, l):
# def log_logical_mesh(mesh: Mesh, *, logger=logging):
#   def _ndig(x):
#     def row(i):
# def tree_global_shape(tree: PyTree, axis_resources: PyTree,
#                       mesh: Mesh) -> PyTree:
# def with_sharding_constraint(x: PyTree, partition_spec: PartitionSpec):


import functools
import itertools
import logging
import re
import jax
import jax.numpy as jnp
import numpy as np
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from vmoe import partitioning
    def mk_dev(process_index):
      return _make_device(process_index=process_index)
    devices = np.asarray([
        [mk_dev(0), mk_dev(0), mk_dev(1)],
        [mk_dev(0), mk_dev(0), mk_dev(2)],
        [mk_dev(0), mk_dev(0), mk_dev(1)],
    ])
    self.assertEqual(
        partitioning.process_has_contiguous_device_slice(
            devices, process_index), expected)

  @parameterized.named_parameters(
      ('false', [[0, 0, 1], [0, 0, 2], [0, 0, 1]], False),
      ('true', [[0, 0, 0], [0, 0, 0], [1, 1, 1]], True),
  )
  def test_processes_have_contiguous_device_slices(self, devices, expected):
    def mk_dev(process_index):
      return _make_device(process_index=process_index)
    devices = np.asarray(devices)
    devices = np.vectorize(mk_dev, otypes=[np.object])(devices)
    self.assertEqual(
        partitioning.processes_have_contiguous_device_slices(devices),
        expected)

  @parameterized.parameters(('other'), ('tpu'))
  def test_get_auto_logical_mesh(self, platform):
    """Tests that the right auto_logical_mesh is run, based on the platform."""
    hardware_mesh = mock.MagicMock()
    device = _make_device(platform=platform)
    with mock.patch.object(

# Path: vmoe/partitioning.py
# def process_has_contiguous_device_slice(devices: np.ndarray,
#                                         process_index: int) -> bool:
# def processes_have_contiguous_device_slices(devices: np.ndarray) -> bool:
# def get_auto_logical_mesh(
#     num_partitions: int,
#     devices: Optional[Sequence[Device]] = None) -> Mesh:
# def get_auto_logical_mesh_other(num_partitions: int,
#                                 hardware_mesh: np.ndarray) -> Mesh:
# def get_auto_logical_mesh_tpu(num_partitions: int,
#                               hardware_mesh: np.ndarray) -> Mesh:
# def get_device_coords(device: Device) -> Union[TpuCoords, OtherCoords]:
# def get_device_coords_other(device: Device) -> OtherCoords:
# def get_device_coords_tpu(device: Device) -> TpuCoords:
# def get_hardware_mesh_local_shape(
#     local_devices: Optional[Sequence[Device]] = None) -> Tuple[int, ...]:
# def get_hardware_mesh_other(devices: Sequence[Device]) -> np.ndarray:
# def get_hardware_mesh_tpu(devices: Sequence[Device]) -> np.ndarray:
# def get_logical_mesh(partitions: Tuple[int, ...],
#                      hardware_mesh: np.ndarray) -> Mesh:
# def get_logical_mesh_default(
#     partitions: Tuple[int, ...],
#     replicas: Tuple[int, ...],
#     hardware_mesh: np.ndarray) -> Mesh:
# def get_logical_mesh_tile_by_process(
#     partitions: Tuple[int, ...],
#     replicas: Tuple[int, ...],
#     hardware_mesh: np.ndarray,
#     hardware_mesh_local_shape: Tuple[int, ...]) -> Mesh:
#   def fn(p, r, l):
# def log_logical_mesh(mesh: Mesh, *, logger=logging):
#   def _ndig(x):
#     def row(i):
# def tree_global_shape(tree: PyTree, axis_resources: PyTree,
#                       mesh: Mesh) -> PyTree:
# def with_sharding_constraint(x: PyTree, partition_spec: PartitionSpec):


# Path: vmoe/partitioning_test.py
import functools
import itertools
import logging
import re
import jax
import jax.numpy as jnp
import numpy as np
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from vmoe import partitioning

    def mk_dev(process_index):
      return _make_device(process_index=process_index)
    devices = np.asarray([
        [mk_dev(0), mk_dev(0), mk_dev(1)],
        [mk_dev(0), mk_dev(0), mk_dev(2)],
        [mk_dev(0), mk_dev(0), mk_dev(1)],
    ])
    self.assertEqual(
        partitioning.process_has_contiguous_device_slice(
            devices, process_index), expected)

  @parameterized.named_parameters(
      ('false', [[0, 0, 1], [0, 0, 2], [0, 0, 1]], False),
      ('true', [[0, 0, 0], [0, 0, 0], [1, 1, 1]], True),
  )
  def test_processes_have_contiguous_device_slices(self, devices, expected):
    def mk_dev(process_index):
      return _make_device(process_index=process_index)
    devices = np.asarray(devices)
    devices = np.vectorize(mk_dev, otypes=[np.object])(devices)
    self.assertEqual(
        partitioning.processes_have_contiguous_device_slices(devices),
        expected)

  @parameterized.parameters(('other'), ('tpu'))
  def test_get_auto_logical_mesh(self, platform):
    """Tests that the right auto_logical_mesh is run, based on the platform."""
    hardware_mesh = mock.MagicMock()
    device = _make_device(platform=platform)
    with mock.patch.object(

        partitioning,