def __init__(self, output=True):
    """Layer remembering the input in forward pass. For reversible models.

    During the first pass through the model this layer saves the input as
    state, and returns the input unmodified. During the second pass through the
    model the layer outputs the input from the first pass. This is used to
    combat numerical stability problems in Terraformer. It doesn't do anything
    in non-reversible models.

    Args:
      output: Whether to pass the input or not.
    """
    n_out = 1 if output else 0
    self._output = output
    super().__init__(name='_RememberInReverse', n_out=n_out)
Layer remembering the input in forward pass. For reversible models.

    During the first pass through the model this layer saves the input as
    state, and returns the input unmodified. During the second pass through the
    model the layer outputs the input from the first pass. This is used to
    combat numerical stability problems in Terraformer. It doesn't do anything
    in non-reversible models.

    Args:
      output: Whether to pass the input or not.
    
trax/layers/research/sparsity.py