def test_types_as_expected(self):
    """Test whether updated metrics match expected types."""

    metric_dict = {
        'cat1': {
            'key': jnp.asarray([1], dtype=jnp.int32),
            'denominator': jnp.asarray([1], dtype=jnp.int16)
        },
        'cat2': {
            'key': 2.0,
            'denominator': jnp.asarray([1], dtype=jnp.bfloat16)
        },
    }
    processed_metrics = metric_utils.update_metrics_dtype(metric_dict)
    self.assertEqual(processed_metrics['cat1']['key'].dtype, jnp.float32)
    self.assertEqual(processed_metrics['cat1']['denominator'].dtype,
                     jnp.float32)
    self.assertIsInstance(processed_metrics['cat2']['key'], float)
    self.assertEqual(processed_metrics['cat2']['denominator'].dtype,
                     jnp.float32)
Test whether updated metrics match expected types.
language/mentionmemory/utils/metric_utils_test.py