docs: def _get_fans(shape):
      r"""Returns the size of input dimension and output dimension, given `shape`.
      
      Args:
        shape: A list of integers.
      
      Returns:
        fan_in: An int. The value of input dimension.
        fan_out: An int. The value of output dimension.
      """
      if len(shape) == 2:
          fan_in = shape[0]
          fan_out = shape[1]
      elif len(shape) == 4 or len(shape) == 5:
          # assuming convolution kernels (2D or 3D).
          kernel_size = np.prod(shape[:2])
          fan_in = shape[-2] * kernel_size
          fan_out = shape[-1] * kernel_size
      else:
          # no specific assumptions
          fan_in = np.sqrt(np.prod(shape))
          fan_out = np.sqrt(np.prod(shape))
      return fan_in, fan_out