# Path: fedjax/core/metrics.py
# class Stat(metaclass=abc.ABCMeta):
# class MeanStat(Stat):
# class SumStat(Stat):
# class Metric(metaclass=abc.ABCMeta):
# class CrossEntropyLoss(Metric):
# class Accuracy(Metric):
# class TopKAccuracy(Metric):
# class SequenceTokenCrossEntropyLoss(Metric):
# class SequenceCrossEntropyLoss(Metric):
# class SequenceTokenAccuracy(Metric):
# class SequenceTokenTopKAccuracy(Metric):
# class SequenceTokenCount(Metric):
# class SequenceCount(Metric):
# class SequenceTruncationRate(Metric):
# class SequenceTokenOOVRate(Metric):
# class SequenceLength(Metric):
# class PerDomainMetric(Metric):
# class ConfusionMatrix(Metric):
#   def result(self) -> jnp.ndarray:
#   def merge(self, other: 'Stat') -> 'Stat':
#   def reduce(self, axis: Optional[int] = 0) -> 'Stat':
#   def __str__(self) -> str:
#   def new(cls, accum, weight) -> 'MeanStat':
#   def result(self) -> jnp.ndarray:
#   def merge(self, other: 'MeanStat') -> 'MeanStat':
#   def reduce(self, axis: Optional[int] = 0) -> 'MeanStat':
#   def new(cls, accum: jnp.ndarray) -> 'SumStat':
#   def result(self) -> jnp.ndarray:
#   def merge(self, other: 'SumStat') -> 'SumStat':
#   def reduce(self, axis: Optional[int] = 0) -> 'SumStat':
#   def zero(self) -> Stat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> Stat:
# def apply_mask(mask: jnp.ndarray, a: jnp.ndarray,
#                b: jnp.ndarray) -> jnp.ndarray:
# def evaluate_batch(metric: Metric,
#                    batch_example: BatchExample,
#                    batch_prediction: BatchPrediction,
#                    batch_mask: Optional[jnp.ndarray] = None) -> Stat:
# def unreduced_cross_entropy_loss(targets: jnp.ndarray,
#                                  preds: jnp.ndarray) -> jnp.ndarray:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
# def get_target_weight(target: jnp.ndarray,
#                       masked_target_values: Tuple[int, ...]) -> jnp.ndarray:
#   def zero(self) -> Stat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> Stat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> SumStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> SumStat:
#   def zero(self) -> SumStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> SumStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> Stat:
#     def broadcast_to(x):
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> Stat:
#     def where(a, b):
#   def zero(self) -> SumStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> SumStat:


from absl.testing import absltest
from absl.testing import parameterized
from fedjax.core import metrics
import jax.numpy as jnp
import numpy.testing as npt
    npt.assert_array_equal(stat.accum, [0., 0., 0., 0.])
    npt.assert_array_equal(stat.weight, [0., 0., 1., 0.])

  @parameterized.named_parameters(
      {
          'testcase_name': 'correct',
          'target': 1,
          'prediction': [0., 1., 0.],
          'expected_result': [[0., 0., 0.], [0., 1., 0.], [0., 0., 0.]],
      }, {
          'testcase_name': 'incorrect',
          'target': 1,
          'prediction': [1., 0., 0.],
          'expected_result': [[0., 0., 0.], [1., 0., 0.], [0., 0., 0.]],
      })
  def test_confusion_matrix(self, target, prediction, expected_result):
    example = {'y': jnp.array(target)}
    prediction = jnp.array(prediction)
    num_classes = len(prediction)
    metric = metrics.ConfusionMatrix(num_classes=num_classes)
    with self.subTest('zero'):
      zero = metric.zero()
      npt.assert_array_equal(zero.accum, jnp.zeros((num_classes, num_classes)))
    with self.subTest('evaluate_example'):
      confusion_matrix = metric.evaluate_example(example, prediction)
      npt.assert_array_equal(confusion_matrix.result(), expected_result)

  def test_confusion_matrix_invalid_input(self):
    example = {'y': jnp.array(1)}
    prediction = jnp.array([0., 1.])

# Path: fedjax/core/metrics.py
# class Stat(metaclass=abc.ABCMeta):
# class MeanStat(Stat):
# class SumStat(Stat):
# class Metric(metaclass=abc.ABCMeta):
# class CrossEntropyLoss(Metric):
# class Accuracy(Metric):
# class TopKAccuracy(Metric):
# class SequenceTokenCrossEntropyLoss(Metric):
# class SequenceCrossEntropyLoss(Metric):
# class SequenceTokenAccuracy(Metric):
# class SequenceTokenTopKAccuracy(Metric):
# class SequenceTokenCount(Metric):
# class SequenceCount(Metric):
# class SequenceTruncationRate(Metric):
# class SequenceTokenOOVRate(Metric):
# class SequenceLength(Metric):
# class PerDomainMetric(Metric):
# class ConfusionMatrix(Metric):
#   def result(self) -> jnp.ndarray:
#   def merge(self, other: 'Stat') -> 'Stat':
#   def reduce(self, axis: Optional[int] = 0) -> 'Stat':
#   def __str__(self) -> str:
#   def new(cls, accum, weight) -> 'MeanStat':
#   def result(self) -> jnp.ndarray:
#   def merge(self, other: 'MeanStat') -> 'MeanStat':
#   def reduce(self, axis: Optional[int] = 0) -> 'MeanStat':
#   def new(cls, accum: jnp.ndarray) -> 'SumStat':
#   def result(self) -> jnp.ndarray:
#   def merge(self, other: 'SumStat') -> 'SumStat':
#   def reduce(self, axis: Optional[int] = 0) -> 'SumStat':
#   def zero(self) -> Stat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> Stat:
# def apply_mask(mask: jnp.ndarray, a: jnp.ndarray,
#                b: jnp.ndarray) -> jnp.ndarray:
# def evaluate_batch(metric: Metric,
#                    batch_example: BatchExample,
#                    batch_prediction: BatchPrediction,
#                    batch_mask: Optional[jnp.ndarray] = None) -> Stat:
# def unreduced_cross_entropy_loss(targets: jnp.ndarray,
#                                  preds: jnp.ndarray) -> jnp.ndarray:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
# def get_target_weight(target: jnp.ndarray,
#                       masked_target_values: Tuple[int, ...]) -> jnp.ndarray:
#   def zero(self) -> Stat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> Stat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> SumStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> SumStat:
#   def zero(self) -> SumStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> SumStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> MeanStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> MeanStat:
#   def zero(self) -> Stat:
#     def broadcast_to(x):
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> Stat:
#     def where(a, b):
#   def zero(self) -> SumStat:
#   def evaluate_example(self, example: SingleExample,
#                        prediction: SinglePrediction) -> SumStat:


# Path: fedjax/core/metrics_test.py
from absl.testing import absltest
from absl.testing import parameterized
from fedjax.core import metrics
import jax.numpy as jnp
import numpy.testing as npt

    npt.assert_array_equal(stat.accum, [0., 0., 0., 0.])
    npt.assert_array_equal(stat.weight, [0., 0., 1., 0.])

  @parameterized.named_parameters(
      {
          'testcase_name': 'correct',
          'target': 1,
          'prediction': [0., 1., 0.],
          'expected_result': [[0., 0., 0.], [0., 1., 0.], [0., 0., 0.]],
      }, {
          'testcase_name': 'incorrect',
          'target': 1,
          'prediction': [1., 0., 0.],
          'expected_result': [[0., 0., 0.], [1., 0., 0.], [0., 0., 0.]],
      })
  def test_confusion_matrix(self, target, prediction, expected_result):
    example = {'y': jnp.array(target)}
    prediction = jnp.array(prediction)
    num_classes = len(prediction)
    metric = metrics.ConfusionMatrix(num_classes=num_classes)
    with self.subTest('zero'):
      zero = metric.zero()
      npt.assert_array_equal(zero.accum, jnp.zeros((num_classes, num_classes)))
    with self.subTest('evaluate_example'):
      confusion_matrix = metric.evaluate_example(example, prediction)
      npt.assert_array_equal(confusion_matrix.result(), expected_result)

  def test_confusion_matrix_invalid_input(self):
    example = {'y': jnp.array(1)}
    prediction = jnp.array([0., 1.])

    metric = metrics.ConfusionMatrix(num_classes=3)