code: def identity(shape: Tuple[int, ...], gain: float = 1) -> JaxArray:
      """Returns the identity matrix. This initializer was proposed in
      `A Simple Way to Initialize Recurrent Networks of Rectified Linear Units
      <https://arxiv.org/abs/1504.00941>`_.
  
      Args:
          shape: Shape of the tensor. It should have exactly rank 2.
          gain: optional scaling factor.
  
      Returns:
          Tensor initialized to the identity matrix.
      """
      assert len(shape) == 2
      return gain * jn.eye(*shape)
  
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 2 for calling a global function with example usage: identity(shape, gain) and returns: unknown