Change the count type from int32 to float32 to avoid overflows.

PiperOrigin-RevId: 723026512

Author: CLU Authors

clu/metrics.py

@@ -784,20 +784,21 @@ class Average(Metric):
 
   @classmethod
   def empty(cls) -> Average:
-    return cls(total=jnp.array(0, jnp.float32), count=jnp.array(0, jnp.int32))
+    return cls(total=jnp.array(0, jnp.float32), count=jnp.array(0, jnp.float32))
 
   @classmethod
   def from_model_output(
       cls, values: jnp.ndarray, mask: jnp.ndarray | None = None, **_
   ) -> Average:
     values, mask = _broadcast_masks(values, mask)
     return cls(
-        total=jnp.where(mask, values, jnp.zeros_like(values)).sum(),
+        total=jnp.where(mask, values, jnp.zeros_like(values)).sum().astype(
+            jnp.float32),
         count=jnp.where(
             mask,
             jnp.ones_like(values, dtype=jnp.int32),
             jnp.zeros_like(values, dtype=jnp.int32),
-        ).sum(),
+        ).sum().astype(jnp.float32),
     )
 
   def merge(self, other: Average) -> Average: