code: def truncated_normal_initializer():
      """TruncatedNormal(0.02) initializer from BERT."""
  
      def init(key, shape, dtype=jnp.float32):
          dtype = jax.dtypes.canonicalize_dtype(dtype)
          return (
              jax.random.truncated_normal(key, -2, 2, shape, dtype)
              * 0.02
          )
  
      return init
  
imports: from typing import Any, Callable, Optional, Sequence, Tuple
  from absl import logging
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 1 for calling a global function with example usage: truncated_normal_initializer() and returns: init, unknown