# 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 vmoe.nn import routing
import chex
import jax
import jax.numpy as jnp
  # parameters of the dispatcher during testing. The output of the
  # NoisyTopExpertsPerItemRouter is supposed to be a dispatcher, but we will
  # simply return the `gates_softmax`, which is fine for testing purposes.
  @mock.patch.object(
      routing.vmoe.moe,
      'get_top_experts_per_item_dispatcher',
      side_effect=lambda x, **_: x)
  def test_forward_deterministic(self, unused_mock):
    """Tests that output is the same given two different gating PRNG seeds."""
    x = jnp.arange(5 * 4).reshape(1, 5, 4).astype(jnp.float32)
    variables = {'params': {'dense': {'kernel': jnp.eye(4)}}}
    layer = routing.NoisyTopExpertsPerItemRouter(
        num_experts=4,
        num_selected_experts=2,
        noise_std=1.0,
        deterministic=True)
    # 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)})
    chex.assert_trees_all_close(y1, y2)
    chex.assert_trees_all_close(m1, m2)

  @mock.patch.object(
      routing.vmoe.moe,
      'get_top_experts_per_item_dispatcher',
      side_effect=lambda x, **_: x)
  def test_forward_not_deterministic(self, unused_mock):
    """Tests that output is different given two different gating PRNG seeds."""
    x = jnp.arange(5 * 4).reshape(1, 5, 4).astype(jnp.float32)
    variables = {'params': {'dense': {'kernel': jnp.eye(4)}}}

# 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/routing_test.py
from unittest import mock
from absl.testing import absltest
from vmoe.nn import routing
import chex
import jax
import jax.numpy as jnp

  # parameters of the dispatcher during testing. The output of the
  # NoisyTopExpertsPerItemRouter is supposed to be a dispatcher, but we will
  # simply return the `gates_softmax`, which is fine for testing purposes.
  @mock.patch.object(
      routing.vmoe.moe,
      'get_top_experts_per_item_dispatcher',
      side_effect=lambda x, **_: x)
  def test_forward_deterministic(self, unused_mock):
    """Tests that output is the same given two different gating PRNG seeds."""
    x = jnp.arange(5 * 4).reshape(1, 5, 4).astype(jnp.float32)
    variables = {'params': {'dense': {'kernel': jnp.eye(4)}}}
    layer = routing.NoisyTopExpertsPerItemRouter(
        num_experts=4,
        num_selected_experts=2,
        noise_std=1.0,
        deterministic=True)
    # 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)})
    chex.assert_trees_all_close(y1, y2)
    chex.assert_trees_all_close(m1, m2)

  @mock.patch.object(
      routing.vmoe.moe,
      'get_top_experts_per_item_dispatcher',
      side_effect=lambda x, **_: x)
  def test_forward_not_deterministic(self, unused_mock):
    """Tests that output is different given two different gating PRNG seeds."""
    x = jnp.arange(5 * 4).reshape(1, 5, 4).astype(jnp.float32)
    variables = {'params': {'dense': {'kernel': jnp.eye(4)}}}

    layer = routing.NoisyTopExpertsPerItemRouter(