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, create_scale, create_offset, use_fast_variance):
      """For all configurations, ensure bf16 outputs from bf16 inputs."""
  
      def f(x):
          ln = layer_norm.LayerNorm(
              axis=-1,
              create_scale=create_scale,
              create_offset=create_offset,
              use_fast_variance=use_fast_variance,
              param_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: import functools
  import itertools
  from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import initializers
  from haiku._src import layer_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 LayerNormTest obj's underlying member method code with example usage: obj.test_bf16(create_scale, create_offset, use_fast_variance) and returns: ln