code: def ResidualZero(*layers, shortcut=None):
      """Wraps a series of layers with a ReZero-style residual connection.
  
      Instead of computing `(shortcut) + (output of layers)`, like in classical
      Residual connection, ResidualZero computes
      `(shortcut) + alpha * (output of layers)`, where `alpha` is a learnable scalar
      initialized with zero.
  
      Args:
        *layers: One or more layers, to be applied in series.
        shortcut: If None (the usual case), the Residual layer computes the
            element-wise sum of the stack-top input with the output of the layer
            series. If specified, the `shortcut` layer applies to a copy of the
            inputs and (elementwise) adds its output to the output from the main
            layer series.
  
      Returns:
          A layer representing a residual connection paired with a layer series.
      """
      layers = _ensure_flat(layers)
      layer = layers[0] if len(layers) == 1 else tl.Serial(layers)
      # TODO(jaszczur): perhaps change inner Serial to Branch?
      return tl.Serial(
          tl.Branch(
              shortcut,
              tl.Serial(
                  layer,
                  tl.Weights(
                      lambda shape, rng: jnp.zeros(
                          shape, dtype=jnp.float32
                      )
                  ),
                  tl.Multiply(),
              ),
          ),
          tl.Add(),  # pylint: disable=no-value-for-parameter
      )
  
file_path: trax/trax/models/research/rezero.py
imports: from trax import layers
  from trax.fastmath import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: ResidualZero() and returns: tl