# Path: vmoe/evaluate/ensemble.py
# def _multiply_no_nan(x: Array, y: Array) -> Array:
# def _get_log_ensemble_softmax(logits: Array, ensemble_size: int) -> Array:
# def ensemble_softmax_xent(logits: Array, labels: Array,
#                           ensemble_size: int) -> Array:
# def label_pred_ensemble_softmax(logits: Array, ensemble_size: int) -> Array:


from absl.testing import absltest
from absl.testing import parameterized
from vmoe.evaluate import ensemble
import jax
import jax.numpy as jnp
import numpy as np
    logits_2d_with_repeat = jnp.asarray(logits_2d_with_repeat)

    labels = jnp.asarray([i % num_classes for i in range(batch_size)])
    one_hot_labels = jax.nn.one_hot(labels, num_classes)

    xent = ensemble.ensemble_softmax_xent(logits_2d_with_repeat,
                                          one_hot_labels, ensemble_size)

    ensemble_softmax = jnp.mean(jax.nn.softmax(logits_3d), 0)
    expected_xent = jnp.asarray([
        -jnp.log(ensemble_softmax[i, labels[i]]) for i in range(batch_size)
    ])

    self.assertSequenceAlmostEqual(list(expected_xent), list(xent), places=5)

  @parameterized.parameters(1, 3, 5)
  def test_label_pred_ensemble_softmax(self, ensemble_size):
    batch_size, num_classes = 16, 8
    logits_3d_shape = (ensemble_size, batch_size, num_classes)
    logits_2d_shape = (ensemble_size * batch_size, num_classes)

    key = jax.random.PRNGKey(666)
    logits_3d = jax.random.normal(key, logits_3d_shape)
    # The ensemble softmax CE assumes the logits have a structure corresponding
    # to jnp.repeat(..., ensemble_size).
    logits_2d_with_repeat = np.zeros(logits_2d_shape)
    for e in range(ensemble_size):
      logits_2d_with_repeat[e::ensemble_size] = logits_3d[e]
    logits_2d_with_repeat = jnp.asarray(logits_2d_with_repeat)


# Path: vmoe/evaluate/ensemble.py
# def _multiply_no_nan(x: Array, y: Array) -> Array:
# def _get_log_ensemble_softmax(logits: Array, ensemble_size: int) -> Array:
# def ensemble_softmax_xent(logits: Array, labels: Array,
#                           ensemble_size: int) -> Array:
# def label_pred_ensemble_softmax(logits: Array, ensemble_size: int) -> Array:


# Path: vmoe/evaluate/ensemble_test.py
from absl.testing import absltest
from absl.testing import parameterized
from vmoe.evaluate import ensemble
import jax
import jax.numpy as jnp
import numpy as np

    logits_2d_with_repeat = jnp.asarray(logits_2d_with_repeat)

    labels = jnp.asarray([i % num_classes for i in range(batch_size)])
    one_hot_labels = jax.nn.one_hot(labels, num_classes)

    xent = ensemble.ensemble_softmax_xent(logits_2d_with_repeat,
                                          one_hot_labels, ensemble_size)

    ensemble_softmax = jnp.mean(jax.nn.softmax(logits_3d), 0)
    expected_xent = jnp.asarray([
        -jnp.log(ensemble_softmax[i, labels[i]]) for i in range(batch_size)
    ])

    self.assertSequenceAlmostEqual(list(expected_xent), list(xent), places=5)

  @parameterized.parameters(1, 3, 5)
  def test_label_pred_ensemble_softmax(self, ensemble_size):
    batch_size, num_classes = 16, 8
    logits_3d_shape = (ensemble_size, batch_size, num_classes)
    logits_2d_shape = (ensemble_size * batch_size, num_classes)

    key = jax.random.PRNGKey(666)
    logits_3d = jax.random.normal(key, logits_3d_shape)
    # The ensemble softmax CE assumes the logits have a structure corresponding
    # to jnp.repeat(..., ensemble_size).
    logits_2d_with_repeat = np.zeros(logits_2d_shape)
    for e in range(ensemble_size):
      logits_2d_with_repeat[e::ensemble_size] = logits_3d[e]
    logits_2d_with_repeat = jnp.asarray(logits_2d_with_repeat)


    labels = ensemble.label_pred_ensemble_softmax(logits_2d_with_repeat,