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

PiperOrigin-RevId: 723026512

Author: CLU Authors

clu/metrics.py

@@ -784,7 +784,7 @@ 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(