def update_value_dtype(value: Any) -> Any:
  """Convert value to more precise type."""
  if isinstance(value, jnp.ndarray):
    return value.astype(jnp.float32)
  return value
Convert value to more precise type.
language/mentionmemory/utils/metric_utils.py