# Path: vmoe/nn/ensemble_routing.py
# def reshape_from_diag_blocks(diagonal_blocks: Array) -> Array:
#   def __call__(self, inputs: Array) -> Tuple[BaseDispatcher, Metrics]:
#   def _compute_gates_softmax_and_metrics(
#       self, inputs: Array) -> Tuple[Array, Metrics]:
#   def _super_compute_gates_softmax_and_metrics(
#       self, inputs: Array) -> Tuple[Array, Metrics]:
# class NoisyTopExpertsPerItemEnsembleRouter(
#     routing.NoisyTopExpertsPerItemRouter):
# 
# Path: vmoe/nn/routing.py
# class NoisyTopExpertsPerItemRouter(nn.Module):
#   def __call__(self, inputs: Array) -> Tuple[BaseDispatcher, Metrics]:
#   def _compute_gates_softmax_and_metrics(
#       self, inputs: Array, num_experts: int) -> Tuple[Array, Metrics]:
#   def _create_dispatcher(self, gates_softmax):
#   def _importance_auxiliary_loss(cls, gates: Array) -> Array:
#   def _load_auxiliary_loss(cls, logits: Array, logits_noisy: Array,
#                            noise_std: Array,
#                            num_selected_experts: int) -> Array:


from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from vmoe.nn import ensemble_routing as ens_routing
from vmoe.nn import routing
import chex
import jax
import jax.numpy as jnp

    selected_experts_no_partitioning = jnp.argmax(
        jnp.dot(x, kernel.reshape(dim_x, num_experts)), -1)
    # Only one group, hence the first 0 index.
    self.assertNotEqual(selected_experts_no_partitioning[0, 0],
                        selected_experts[0, 0])
    self.assertNotEqual(selected_experts_no_partitioning[0, 1],
                        selected_experts[0, 1])

  def test_diag_blocks(self):
    """Test the utils to manipulate the block-diagonal representations."""
    rows = 2
    cols = 3
    ens_size = 4

    diag_blocks1 = jnp.asarray(
        [k * jnp.ones((rows, cols)) for k in range(1, ens_size + 1)])
    full_block_diag_matrix1 = jax.scipy.linalg.block_diag(*diag_blocks1)

    diag_blocks2 = jnp.asarray(
        [-k * jnp.ones((rows, cols)) for k in range(1, ens_size + 1)])
    full_block_diag_matrix2 = jax.scipy.linalg.block_diag(*diag_blocks2)

    diag_blocks12 = jnp.asarray([diag_blocks1, diag_blocks2])
    full_block_diag_matrix12 = jnp.asarray(
        [full_block_diag_matrix1, full_block_diag_matrix2])

    def dist(u, v):
      return float(jnp.sum(jnp.abs(u - v)))


# Path: vmoe/nn/ensemble_routing.py
# def reshape_from_diag_blocks(diagonal_blocks: Array) -> Array:
#   def __call__(self, inputs: Array) -> Tuple[BaseDispatcher, Metrics]:
#   def _compute_gates_softmax_and_metrics(
#       self, inputs: Array) -> Tuple[Array, Metrics]:
#   def _super_compute_gates_softmax_and_metrics(
#       self, inputs: Array) -> Tuple[Array, Metrics]:
# class NoisyTopExpertsPerItemEnsembleRouter(
#     routing.NoisyTopExpertsPerItemRouter):
# 
# Path: vmoe/nn/routing.py
# class NoisyTopExpertsPerItemRouter(nn.Module):
#   def __call__(self, inputs: Array) -> Tuple[BaseDispatcher, Metrics]:
#   def _compute_gates_softmax_and_metrics(
#       self, inputs: Array, num_experts: int) -> Tuple[Array, Metrics]:
#   def _create_dispatcher(self, gates_softmax):
#   def _importance_auxiliary_loss(cls, gates: Array) -> Array:
#   def _load_auxiliary_loss(cls, logits: Array, logits_noisy: Array,
#                            noise_std: Array,
#                            num_selected_experts: int) -> Array:


# Path: vmoe/nn/ensemble_routing_test.py
from unittest import mock
from absl.testing import absltest
from absl.testing import parameterized
from vmoe.nn import ensemble_routing as ens_routing
from vmoe.nn import routing
import chex
import jax
import jax.numpy as jnp


    selected_experts_no_partitioning = jnp.argmax(
        jnp.dot(x, kernel.reshape(dim_x, num_experts)), -1)
    # Only one group, hence the first 0 index.
    self.assertNotEqual(selected_experts_no_partitioning[0, 0],
                        selected_experts[0, 0])
    self.assertNotEqual(selected_experts_no_partitioning[0, 1],
                        selected_experts[0, 1])

  def test_diag_blocks(self):
    """Test the utils to manipulate the block-diagonal representations."""
    rows = 2
    cols = 3
    ens_size = 4

    diag_blocks1 = jnp.asarray(
        [k * jnp.ones((rows, cols)) for k in range(1, ens_size + 1)])
    full_block_diag_matrix1 = jax.scipy.linalg.block_diag(*diag_blocks1)

    diag_blocks2 = jnp.asarray(
        [-k * jnp.ones((rows, cols)) for k in range(1, ens_size + 1)])
    full_block_diag_matrix2 = jax.scipy.linalg.block_diag(*diag_blocks2)

    diag_blocks12 = jnp.asarray([diag_blocks1, diag_blocks2])
    full_block_diag_matrix12 = jnp.asarray(
        [full_block_diag_matrix1, full_block_diag_matrix2])

    def dist(u, v):
      return float(jnp.sum(jnp.abs(u - v)))


    out12 = ens_routing.reshape_from_diag_blocks(diag_blocks12)