code: def orthogonal(
      shape: Tuple[int, ...], gain: float = 1, axis: int = -1
  ) -> JaxArray:
      """Returns a uniformly distributed orthogonal tensor from
      `Exact solutions to the nonlinear dynamics of learning in deep linear neural networks
      <https://openreview.net/forum?id=_wzZwKpTDF_9C>`_.
  
      Args:
          shape: shape of the output tensor.
          gain: optional scaling factor.
          axis: the orthogonalizarion axis
  
      Returns:
          An orthogonally initialized tensor.
          These tensors will be row-orthonormal along the access specified by
          ``axis``. If the rank of the weight is greater than 2, the shape will be
          flattened in all other dimensions and then will be row-orthonormal along the
          final dimension. Note that this only works if the ``axis`` dimension is
          larger, otherwise the tensor will be transposed (equivalently, it will be
          column orthonormal instead of row orthonormal).
          If the shape is not square, the matrices will have orthonormal rows or
          columns depending on which side is smaller.
      """
      n_rows = shape[axis]
      n_cols = np.prod(shape) // n_rows
      matrix_shape = (
          (n_rows, n_cols) if n_rows > n_cols else (n_cols, n_rows)
      )
      norm_dst = random.normal(matrix_shape)
      q_mat, r_mat = np.linalg.qr(norm_dst)
      # Enforce Q is uniformly distributed
      q_mat *= np.sign(np.diag(r_mat))
      if n_rows < n_cols:
          q_mat = q_mat.T
      q_mat = np.reshape(
          q_mat, (n_rows,) + tuple(np.delete(shape, axis))
      )
      q_mat = np.moveaxis(q_mat, 0, axis)
      return gain * jn.array(q_mat)
  
file_path: objax/objax/nn/init.py
imports: from typing import Tuple
  import numpy
  import scipy.stats
  from jax import numpy
  from objax import random
  from objax.typing import JaxArray
label_desc: Use image node_id 6 for calling a global function with example usage: orthogonal(shape, gain, axis) and returns: unknown