# Path: t5x/metrics.py
# def _check_param(value, *, ndim=None, dtype=jnp.float32):
#   def from_model_output(cls, values: Scalar, **_) -> clu_metrics.Metric:
#   def merge(self, other: "Sum") -> "Sum":
#   def compute(self) -> Scalar:
#   def replace_steps(self, steps) -> "Step":
#   def compute(self) -> Scalar:
#   def from_model_output(cls,
#                         values: Scalar,
#                         steps: Optional[int] = 1,
#                         **_) -> clu_metrics.Metric:
#   def merge(self, other: "AveragePerStep") -> "AveragePerStep":
#   def compute(self) -> Scalar:
#   def merge(self, other: "Time") -> "Time":
#   def compute(self) -> Scalar:
#   def replace_duration(self, duration: Scalar) -> "Time":
#   def from_model_output(cls, numerator: float, **_) -> clu_metrics.Metric:
#   def merge(self, other: "TimeRate") -> "TimeRate":
#   def compute(self) -> Scalar:
#   def replace_duration(self, duration: Scalar) -> "Time":
#   def from_model_output(cls,
#                         steps: Optional[int] = 1,
#                         **_) -> clu_metrics.Metric:
#   def merge(self, other: "StepsPerTime") -> "StepsPerTime":
#   def compute(self) -> Scalar:
# def is_metric_obj(obj):
# def is_time_metric(obj):
# def create_metrics_dict(float_metrics_dict):
# def shape_obj_to_defined_obj(obj: clu_metrics.Metric):
#   def class_attr_shape(a):
# def set_time_metrics_duration(metrics, duration):
#   def fn(o):
# def set_step_metrics_num_steps(metrics, num_steps):
#   def fn(o):
# class Sum(clu_metrics.Metric):
# class Step(clu_metrics.Metric):
# class AveragePerStep(Step):
# class Time(clu_metrics.Metric):
# class TimeRate(Time):
# class StepsPerTime(Step, Time):


from absl.testing import absltest
from absl.testing import parameterized
from t5x import metrics
import jax
import jax.numpy as jnp
import numpy as np
# Copyright 2022 The T5X Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for clu.metrics."""



class MetricsTest(parameterized.TestCase):

  @parameterized.named_parameters(
      ("0d_values", 2., 2.), ("1d_values", [1, 2, 3], 6.),
      ("2d_values", [[1, 2], [2, 3], [3, 4]], 15.),
      ("3d_values", [[[1, 2], [2, 3]], [[2, 1], [3, 4]], [[3, 1], [4, 1]]], 27.)
  )
  def test_sum(self, values, expected_result):
    self.assertAlmostEqual(

# Path: t5x/metrics.py
# def _check_param(value, *, ndim=None, dtype=jnp.float32):
#   def from_model_output(cls, values: Scalar, **_) -> clu_metrics.Metric:
#   def merge(self, other: "Sum") -> "Sum":
#   def compute(self) -> Scalar:
#   def replace_steps(self, steps) -> "Step":
#   def compute(self) -> Scalar:
#   def from_model_output(cls,
#                         values: Scalar,
#                         steps: Optional[int] = 1,
#                         **_) -> clu_metrics.Metric:
#   def merge(self, other: "AveragePerStep") -> "AveragePerStep":
#   def compute(self) -> Scalar:
#   def merge(self, other: "Time") -> "Time":
#   def compute(self) -> Scalar:
#   def replace_duration(self, duration: Scalar) -> "Time":
#   def from_model_output(cls, numerator: float, **_) -> clu_metrics.Metric:
#   def merge(self, other: "TimeRate") -> "TimeRate":
#   def compute(self) -> Scalar:
#   def replace_duration(self, duration: Scalar) -> "Time":
#   def from_model_output(cls,
#                         steps: Optional[int] = 1,
#                         **_) -> clu_metrics.Metric:
#   def merge(self, other: "StepsPerTime") -> "StepsPerTime":
#   def compute(self) -> Scalar:
# def is_metric_obj(obj):
# def is_time_metric(obj):
# def create_metrics_dict(float_metrics_dict):
# def shape_obj_to_defined_obj(obj: clu_metrics.Metric):
#   def class_attr_shape(a):
# def set_time_metrics_duration(metrics, duration):
#   def fn(o):
# def set_step_metrics_num_steps(metrics, num_steps):
#   def fn(o):
# class Sum(clu_metrics.Metric):
# class Step(clu_metrics.Metric):
# class AveragePerStep(Step):
# class Time(clu_metrics.Metric):
# class TimeRate(Time):
# class StepsPerTime(Step, Time):


# Path: t5x/metrics_test.py
from absl.testing import absltest
from absl.testing import parameterized
from t5x import metrics
import jax
import jax.numpy as jnp
import numpy as np

# Copyright 2022 The T5X Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for clu.metrics."""



class MetricsTest(parameterized.TestCase):

  @parameterized.named_parameters(
      ("0d_values", 2., 2.), ("1d_values", [1, 2, 3], 6.),
      ("2d_values", [[1, 2], [2, 3], [3, 4]], 15.),
      ("3d_values", [[[1, 2], [2, 3]], [[2, 1], [3, 4]], [[3, 1], [4, 1]]], 27.)
  )
  def test_sum(self, values, expected_result):
    self.assertAlmostEqual(

        metrics.Sum.from_model_output(values).compute(), expected_result)