code: def OrthogonalInitializer(stddev=1.0):
      """Returns an orthogonal initializer."""
  
      def Init(shape, rng):
          """Returns orthogonalized random normal values with the given `shape`."""
          # Have at least 2 elements in shape.
          cur_shape = list(shape)
          while len(cur_shape) < 2:
              cur_shape = [1] + cur_shape
  
          # Flatten the input shape with the last dimension remaining.
          n_rows = 1
          for dim in cur_shape[:-1]:
              n_rows *= dim
          n_cols = cur_shape[-1]
          flat_shape = (
              (n_cols, n_rows) if n_rows < n_cols else (n_rows, n_cols)
          )
  
          # Generate a random matrix
          a = random.normal(rng, flat_shape, dtype=jnp.float32)
  
          # Compute the qr factorization
          q, r = jnp.linalg.qr(a)
  
          # Make Q uniform
          d = jnp.diag(r)
          q *= jnp.sign(d)
  
          # Transpose and reshape back q if needed.
          if n_rows < n_cols:
              q = jnp.transpose(q)
          q = jnp.reshape(q, shape)
  
          # Return scaled as requested.
          return stddev * q
  
      return Init
  
file_path: trax/trax/layers/initializers.py
imports: from absl import logging
  import numpy
  import tensorflow.compat.v2
  from trax.fastmath import numpy
  from trax.fastmath import random
label_desc: Use image node_id 13 for calling a global function with example usage: OrthogonalInitializer(stddev) and returns: Init, unknown