class_docstr: Residual connection module.
  
  Attributes:
    residual_type: str; residual connection type. Possible values are [
      'gated', 'sigtanh', 'rezero', 'highway', 'add'].
    dtype: Jax dtype; The dtype of the computation (default: float32).
class_docstr_tok: ["Residual","connection","module",".","Attributes",":","residual_type",":","str",";","residual","connection","type",".","Possible","values","are","[","'gated","'",",","'sigtanh","'",",","'rezero","'",",","'highway","'",",","'add","'","]",".","dtype",":","Jax","dtype",";","The","dtype","of","the","computation","(","default",":","float32",")","."]
docstr: Applies the residual connection on given input/output of a module.
  
  Args:
    x: Input of the module.
    y: Output of the module.
  
  Returns:
    Output: A combination of the x and y.
docstr_tok: ["Applies","the","residual","connection","on","given","input\/output","of","a","module",".","Args",":","x",":","Input","of","the","module",".","y",":","Output","of","the","module",".","Returns",":","Output",":","A","combination","of","the","x","and","y","."]
code: def __call__(self, x, y):
      """Applies the residual connection on given input/output of a module.
  
      Args:
        x: Input of the module.
        y: Output of the module.
  
      Returns:
        Output: A combination of the x and y.
      """
      if x.shape != y.shape:
          raise ValueError("x and y should be of the same shape.")
  
      dtype = self.dtype
  
      if self.residual_type == "add":
          return x + y
  
      elif self.residual_type == "highway":
          features = x.shape[-1]
          hw_gate = nn.sigmoid(
              nn.Dense(
                  features=features,
                  use_bias=True,
                  kernel_init=initializers.zeros,
                  bias_init=lambda rng, shape, *_: jnp.full(
                      shape, -10.0
                  ),
                  dtype=dtype,
              )(x)
          )
          output = jnp.multiply((1 - hw_gate), x) + jnp.multiply(
              hw_gate, y
          )
  
      elif self.residual_type == "rezero":
          # Based on https://arxiv.org/pdf/2003.04887v1.pdf.
          alpha = self.param("rezero_alpha", initializers.zeros, (1,))
          return x + (alpha * y)
  
      elif self.residual_type == "sigtanh":
          # Based on https://arxiv.org/pdf/1606.05328.pdf.
          features = x.shape[-1]
          # sigmoid(W_g.y).
          sigmoid_y = nn.sigmoid(
              nn.Dense(
                  features=features,
                  use_bias=True,
                  kernel_init=initializers.zeros,
                  bias_init=lambda rng, shape, *_: jnp.full(
                      shape, -10.0
                  ),
                  dtype=dtype,
              )(y)
          )
          # tanh(U_g.y).
          tanh_y = nn.tanh(
              nn.Dense(
                  features=features,
                  use_bias=False,
                  kernel_init=initializers.zeros,
                  bias_init=initializers.zeros,
                  dtype=dtype,
              )(y)
          )
          return x + (sigmoid_y * tanh_y)
  
      elif self.residual_type == "gated":
          # Based on https://arxiv.org/pdf/1910.06764.pdf.
          features = x.shape[-1]
          # Reset gate: r = sigmoid(W_r.x + U_r.y).
          r = nn.sigmoid(
              nn.Dense(
                  features=features,
                  use_bias=False,
                  kernel_init=initializers.zeros,
                  bias_init=initializers.zeros,
                  dtype=dtype,
              )(x)
              + nn.Dense(
                  features=features,
                  use_bias=False,
                  kernel_init=initializers.zeros,
                  bias_init=initializers.zeros,
                  dtype=dtype,
              )(y)
          )
          # Update gate: z = sigmoid(W_z.x + U_z.y - b_g).
          # NOTE: the paper claims best initializtion for their task for b is 2.
          b_g = self.param(
              "b_g",
              lambda rng, shape, *_: jnp.full(shape, 10.0),
              (features,),
          ).astype(dtype)
          z = nn.sigmoid(
              nn.Dense(
                  features=features,
                  use_bias=False,
                  kernel_init=initializers.zeros,
                  bias_init=initializers.zeros,
                  dtype=dtype,
              )(x)
              + nn.Dense(
                  features=features,
                  use_bias=False,
                  kernel_init=initializers.zeros,
                  bias_init=initializers.zeros,
                  dtype=dtype,
              )(y)
              - b_g
          )
          # Candidate_activation: h' = tanh(W_g.y + U_g.(r*x)).
          h = jnp.tanh(
              nn.Dense(
                  features=features,
                  use_bias=False,
                  kernel_init=initializers.zeros,
                  bias_init=initializers.zeros,
                  dtype=dtype,
              )(y)
              + nn.Dense(
                  features=features,
                  use_bias=False,
                  kernel_init=initializers.zeros,
                  bias_init=initializers.zeros,
                  dtype=dtype,
              )(jnp.multiply(r, x))
          )
  
          # Output: g = (1-z)*x + z*h.
          output = jnp.multiply((1.0 - z), x) + jnp.multiply(z, h)
  
      else:
          raise ValueError(
              f"Residual type {self.residual_type} is not defined."
          )
      return output
  
imports: from typing import Callable, Iterable, Optional, Sequence
  import flax.linen
  import jax
  from jax.nn import initializers
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the Residual obj's underlying member method code with example usage: obj.__call__(x, y) and returns: output, unknown, unknown, unknown