# Path: vmoe/moe.py
# class BaseDispatcher(abc.ABC):
# class EinsumDispatcher(BaseDispatcher):
# class ExpertIndicesDispatcher(BaseDispatcher):
# class Bfloat16Dispatcher(BaseDispatcher):
#   def dispatch(self, data: Array) -> Array:
#   def combine(self, data: Array) -> Array:
#   def dispatch(self, data: Array) -> Array:
#   def combine(self, data: Array) -> Array:
#   def dispatch(self, data: Array) -> Array:
#   def combine(self, data: Array) -> Array:
#   def dispatch(self, data: Array) -> Array:
#   def combine(self, data: Array) -> Array:
# def get_top_experts_per_item_dispatcher(gates: Array, name: str,
#                                         num_selected_experts: int,
#                                         capacity: int, batch_priority: bool,
#                                         **dispatcher_kwargs) -> BaseDispatcher:
# def sparse_moe_spmd(target: flax.linen.transforms.Target,
#                     split_rngs: bool = False,
#                     has_aux: bool = False,
#                     methods=None):
#   def wrapper(expert_fn: Callable[..., Any]):
#     def transformed(scopes, dispatcher, *inputs):
# def _cast_to_bfloat16(x: Array) -> Array:
# def _convert_partition_spec(spec):
# def _dispatch(data: Array, partition_spec: Optional[PartitionSpec]) -> Array:
# def _receive(data: Array, num_groups: int,
#              partition_spec: Optional[PartitionSpec]) -> Array:
# def _scatter_nd(indices, updates, shape):
# def _get_top_experts_per_item_common(
#     gates: Array, num_selected_experts: int,
#     batch_priority: bool) -> Tuple[Array, Array, Array]:
# def _get_top_experts_per_item_einsum_dispatcher(
#     gates: Array, num_selected_experts: int, capacity: int,
#     batch_priority: bool, **dispatcher_kwargs) -> EinsumDispatcher:
# def _get_top_experts_per_item_expert_indices_dispatcher(
#     gates: Array, num_selected_experts: int, capacity: int,
#     batch_priority: bool, **dispatcher_kwargs) -> ExpertIndicesDispatcher:


from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from vmoe import moe
import flax.core
import flax.linen as nn
import jax
import jax.experimental.maps as maps
import jax.experimental.pjit as pjit
import jax.numpy as jnp
import numpy as np
        gates,
        num_selected_experts=2,
        capacity=2,
        batch_priority=batch_priority)
    np.testing.assert_array_almost_equal(dispatcher.combine_weights, expected)

  @parameterized.named_parameters(
      ('vanilla', False,
       [[(0, 0), (2, 1)], [(0, 1), (2, 2)], [(2, 0), (0, 2)]]),
      ('bpr', True,
       [[(0, 1), (2, 2)], [(0, 0), (2, 1)], [(2, 0), (0, 2)]]),
  )
  def test_top_experts_per_item_expert_indices(self, batch_priority,
                                               expected_indices,
                                               ):
    gates = jnp.asarray([[.4, .0, .2], [.5, .1, .4], [.3, .2, .6]])
    dispatcher = moe._get_top_experts_per_item_expert_indices_dispatcher(
        gates,
        num_selected_experts=2,
        capacity=2,
        batch_priority=batch_priority)
    self.assertEqual(dispatcher.num_experts, 3)
    self.assertEqual(dispatcher.capacity, 2)
    np.testing.assert_array_equal(dispatcher.indices, expected_indices)
    expected_combine_weights = jnp.asarray([[.4, .2], [.5, .4], [.6, .3]])
    np.testing.assert_array_almost_equal(dispatcher.combine_weights,
                                         expected_combine_weights)

  @parameterized.named_parameters(
      ('einsum', 'einsum', moe.EinsumDispatcher),

# Path: vmoe/moe.py
# class BaseDispatcher(abc.ABC):
# class EinsumDispatcher(BaseDispatcher):
# class ExpertIndicesDispatcher(BaseDispatcher):
# class Bfloat16Dispatcher(BaseDispatcher):
#   def dispatch(self, data: Array) -> Array:
#   def combine(self, data: Array) -> Array:
#   def dispatch(self, data: Array) -> Array:
#   def combine(self, data: Array) -> Array:
#   def dispatch(self, data: Array) -> Array:
#   def combine(self, data: Array) -> Array:
#   def dispatch(self, data: Array) -> Array:
#   def combine(self, data: Array) -> Array:
# def get_top_experts_per_item_dispatcher(gates: Array, name: str,
#                                         num_selected_experts: int,
#                                         capacity: int, batch_priority: bool,
#                                         **dispatcher_kwargs) -> BaseDispatcher:
# def sparse_moe_spmd(target: flax.linen.transforms.Target,
#                     split_rngs: bool = False,
#                     has_aux: bool = False,
#                     methods=None):
#   def wrapper(expert_fn: Callable[..., Any]):
#     def transformed(scopes, dispatcher, *inputs):
# def _cast_to_bfloat16(x: Array) -> Array:
# def _convert_partition_spec(spec):
# def _dispatch(data: Array, partition_spec: Optional[PartitionSpec]) -> Array:
# def _receive(data: Array, num_groups: int,
#              partition_spec: Optional[PartitionSpec]) -> Array:
# def _scatter_nd(indices, updates, shape):
# def _get_top_experts_per_item_common(
#     gates: Array, num_selected_experts: int,
#     batch_priority: bool) -> Tuple[Array, Array, Array]:
# def _get_top_experts_per_item_einsum_dispatcher(
#     gates: Array, num_selected_experts: int, capacity: int,
#     batch_priority: bool, **dispatcher_kwargs) -> EinsumDispatcher:
# def _get_top_experts_per_item_expert_indices_dispatcher(
#     gates: Array, num_selected_experts: int, capacity: int,
#     batch_priority: bool, **dispatcher_kwargs) -> ExpertIndicesDispatcher:


# Path: vmoe/moe_test.py
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from vmoe import moe
import flax.core
import flax.linen as nn
import jax
import jax.experimental.maps as maps
import jax.experimental.pjit as pjit
import jax.numpy as jnp
import numpy as np

        gates,
        num_selected_experts=2,
        capacity=2,
        batch_priority=batch_priority)
    np.testing.assert_array_almost_equal(dispatcher.combine_weights, expected)

  @parameterized.named_parameters(
      ('vanilla', False,
       [[(0, 0), (2, 1)], [(0, 1), (2, 2)], [(2, 0), (0, 2)]]),
      ('bpr', True,
       [[(0, 1), (2, 2)], [(0, 0), (2, 1)], [(2, 0), (0, 2)]]),
  )
  def test_top_experts_per_item_expert_indices(self, batch_priority,
                                               expected_indices,
                                               ):
    gates = jnp.asarray([[.4, .0, .2], [.5, .1, .4], [.3, .2, .6]])
    dispatcher = moe._get_top_experts_per_item_expert_indices_dispatcher(
        gates,
        num_selected_experts=2,
        capacity=2,
        batch_priority=batch_priority)
    self.assertEqual(dispatcher.num_experts, 3)
    self.assertEqual(dispatcher.capacity, 2)
    np.testing.assert_array_equal(dispatcher.indices, expected_indices)
    expected_combine_weights = jnp.asarray([[.4, .2], [.5, .4], [.6, .3]])
    np.testing.assert_array_almost_equal(dispatcher.combine_weights,
                                         expected_combine_weights)

  @parameterized.named_parameters(
      ('einsum', 'einsum', moe.EinsumDispatcher),

      ('indices', 'indices', moe.ExpertIndicesDispatcher),