code: def to_bf16(t):
      return jax.tree_map(
          lambda x: x.astype(jnp.bfloat16)
          if x.dtype == jnp.float32
          else x,
          t,
      )
  
imports: import jax
  import jax.numpy
  from jax.experimental.pjit import with_sharding_constraint
  from optax import AdditiveWeightDecayState, GradientTransformation, OptState
label_desc: Use image node_id 7 for calling a global function with example usage: to_bf16(t) and returns: jax