# 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
    with self.assertRaises(ValueError):
      fn()

  def test_time(self):
    duration = 2.
    metric = metrics.Time().replace_duration(duration)
    self.assertAlmostEqual(metric.compute(), duration)

  def test_time_unset_duration(self):
    metric = metrics.Time()
    with self.assertRaises(ValueError):
      metric.compute()

  @parameterized.named_parameters(
      ("0d_values", 2., 2.),
      ("1d_values", [1, 2, 3], 6.),
  )
  def test_average_per_step(self, values, expected_result):
    a = metrics.AveragePerStep.from_model_output(values)
    m = metrics.set_step_metrics_num_steps({"a": a}, 1)
    self.assertAlmostEqual(m["a"].compute(), expected_result)

    steps = 5
    b = metrics.AveragePerStep.from_model_output(values, steps=steps)
    m = metrics.set_step_metrics_num_steps({"b": b}, steps)
    self.assertAlmostEqual(m["b"].compute(), expected_result / steps)

  def test_steps_per_time(self):
    steps = 8.
    duration = 2.

# 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

    with self.assertRaises(ValueError):
      fn()

  def test_time(self):
    duration = 2.
    metric = metrics.Time().replace_duration(duration)
    self.assertAlmostEqual(metric.compute(), duration)

  def test_time_unset_duration(self):
    metric = metrics.Time()
    with self.assertRaises(ValueError):
      metric.compute()

  @parameterized.named_parameters(
      ("0d_values", 2., 2.),
      ("1d_values", [1, 2, 3], 6.),
  )
  def test_average_per_step(self, values, expected_result):
    a = metrics.AveragePerStep.from_model_output(values)
    m = metrics.set_step_metrics_num_steps({"a": a}, 1)
    self.assertAlmostEqual(m["a"].compute(), expected_result)

    steps = 5
    b = metrics.AveragePerStep.from_model_output(values, steps=steps)
    m = metrics.set_step_metrics_num_steps({"b": b}, steps)
    self.assertAlmostEqual(m["b"].compute(), expected_result / steps)

  def test_steps_per_time(self):
    steps = 8.
    duration = 2.

    metric = metrics.StepsPerTime.from_model_output(