Change metric compute() return type to jnp.ndarray to match upcoming …

…pytype changes in CLU.

PiperOrigin-RevId: 557561097

t5x/metrics.py

@@ -25,10 +25,9 @@
 import flax  # Only used for flax.struct.dataclass.
 import jax
 import jax.numpy as jnp
-import numpy as np
 
 MetricsMap = MutableMapping[str, clu_metrics.Metric]
-Scalar = Union[int, float, np.number, np.ndarray, jnp.ndarray]
+Scalar = jnp.ndarray
 
 
 def _check_param(value, *, ndim=None, dtype=jnp.float32):
@@ -91,15 +90,15 @@ class Step(clu_metrics.Metric):
   """
   steps: Optional[int] = 1
 
-  def replace_steps(self, steps) -> "Step":
+  def replace_steps(self, steps: Union[Scalar, int]) -> "Step":
     return self.replace(steps=steps)
 
   def compute(self) -> Scalar:
     if self.steps is None:
       raise ValueError(
           "`steps` must be set by calling `replace_steps` before computing metric."
       )
-    return self.steps
+    return jnp.array(self.steps)
 
 
 @flax.struct.dataclass
@@ -176,7 +175,7 @@ def replace_duration(self, duration: Scalar) -> "Time":
     Returns:
       A new Time object.
     """
-    return self.replace(duration=duration)
+    return self.replace(duration=jnp.array(duration))
 
 
 @flax.struct.dataclass