code: def truncated_normal_initializer(
      stddev: float = 1e-2, dtype: jnp.dtype = jnp.float_
  ) -> Initializer:
      """Returns a truncated normal parameter initializer.
  
      The truncation bounds are -2 and +2 standard deviations.
  
      Args:
        stddev: The standard deviation of the truncated normal distribution.
        dtype: The data type to use.
  
      Returns:
        Initializer function compatible with Flax modules.
      """
  
      def init(key, shape, dtype=dtype):
          dtype = jax.dtypes.canonicalize_dtype(dtype)
          if jnp.issubdtype(dtype, jnp.floating):
              # constant is stddev of standard normal truncated to (-2, 2)
              s = stddev / jnp.array(0.87962566103423978, dtype)
          else:
              # constant is stddev of complex standard normal truncated to 2
              s = stddev / jnp.array(0.95311164380491208, dtype)
          return (
              jax.random.truncated_normal(key, -2, 2, shape, dtype) * s
          )
  
      return init
  
file_path: scenic/scenic/model_lib/layers/nn_ops.py
imports: from typing import Callable, Sequence
  import flax.linen
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 12 for calling a global function with example usage: truncated_normal_initializer(stddev, dtype) and returns: init, unknown