# 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
    batch_size = 6
    batch_size_with_repeat = batch_size * ens_size
    dim_x = num_experts
    product_shape = dim_x * batch_size_with_repeat * num_groups
    x = jnp.arange(product_shape)
    x = x.reshape(num_groups, batch_size_with_repeat, dim_x).astype(jnp.float32)
    kernel = jnp.eye(dim_x).reshape(dim_x, ens_size, num_experts // ens_size)
    variables = {'params': {'dense': {'kernel': kernel}}}
    layer = NoisyTopExpertsPerItemEnsembleRouter(
        num_experts=num_experts,
        ensemble_size=ens_size,
        num_selected_experts=num_selected_experts,
        noise_std=1.0,
        deterministic=False)
    # y's are dispatch weights, m's are metrics.
    y1, m1 = layer.apply(variables, x, rngs={'gating': jax.random.PRNGKey(0)})
    y2, m2 = layer.apply(variables, x, rngs={'gating': jax.random.PRNGKey(1)})
    different_fn = lambda x, y: jnp.abs(x - y).sum() > 1e-3
    error_msg_fn = lambda x, y: f'{x} is too close to {y}'
    chex.assert_trees_all_equal_comparator(different_fn, error_msg_fn, y1, y2)
    # Importance loss is applied before adding noise, so it should be identical.
    chex.assert_trees_all_close(m1['importance_loss'], m2['importance_loss'])
    del m1['importance_loss']
    del m2['importance_loss']
    chex.assert_trees_all_equal_comparator(different_fn, error_msg_fn, m1, m2)
    for loss1, loss2 in zip(m1.values(), m2.values()):
      self.assertEqual(loss1.shape, expected_output_shape)
      self.assertEqual(loss2.shape, expected_output_shape)

  @mock.patch.object(

# 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

    batch_size = 6
    batch_size_with_repeat = batch_size * ens_size
    dim_x = num_experts
    product_shape = dim_x * batch_size_with_repeat * num_groups
    x = jnp.arange(product_shape)
    x = x.reshape(num_groups, batch_size_with_repeat, dim_x).astype(jnp.float32)
    kernel = jnp.eye(dim_x).reshape(dim_x, ens_size, num_experts // ens_size)
    variables = {'params': {'dense': {'kernel': kernel}}}
    layer = NoisyTopExpertsPerItemEnsembleRouter(
        num_experts=num_experts,
        ensemble_size=ens_size,
        num_selected_experts=num_selected_experts,
        noise_std=1.0,
        deterministic=False)
    # y's are dispatch weights, m's are metrics.
    y1, m1 = layer.apply(variables, x, rngs={'gating': jax.random.PRNGKey(0)})
    y2, m2 = layer.apply(variables, x, rngs={'gating': jax.random.PRNGKey(1)})
    different_fn = lambda x, y: jnp.abs(x - y).sum() > 1e-3
    error_msg_fn = lambda x, y: f'{x} is too close to {y}'
    chex.assert_trees_all_equal_comparator(different_fn, error_msg_fn, y1, y2)
    # Importance loss is applied before adding noise, so it should be identical.
    chex.assert_trees_all_close(m1['importance_loss'], m2['importance_loss'])
    del m1['importance_loss']
    del m2['importance_loss']
    chex.assert_trees_all_equal_comparator(different_fn, error_msg_fn, m1, m2)
    for loss1, loss2 in zip(m1.values(), m2.values()):
      self.assertEqual(loss1.shape, expected_output_shape)
      self.assertEqual(loss2.shape, expected_output_shape)

  @mock.patch.object(

      routing.vmoe.moe,