docstr: For all configurations, ensure bf16 outputs from bf16 inputs.
docstr_tok: ["For","all","configurations",",","ensure","bf16","outputs","from","bf16","inputs","."]
code: def test_bf16(self):
      """For all configurations, ensure bf16 outputs from bf16 inputs."""
  
      def f(x):
          ln = rms_norm.RMSNorm(axis=-1)
          return ln(x)
  
      fwd = transform.transform(f)
      data = jnp.zeros([2, 3, 4, 5], dtype=jnp.bfloat16)
      params = fwd.init(jax.random.PRNGKey(428), data)
      bf16_params = jax.tree_util.tree_map(
          lambda t: t.astype(jnp.bfloat16), params
      )
      self.assertEqual(
          fwd.apply(bf16_params, None, data).dtype, jnp.bfloat16
      )
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import initializers
  from haiku._src import rms_norm
  from haiku._src import test_utils
  from haiku._src import transform
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 2 for calling the RMSNormTest obj's underlying member method code with example usage: obj.test_bf16() and returns: ln