class_docstr: Half of a RevNet-style residual that optionally performs attention.
  
  When attention_layer is None, this layer has the signature ::
  
      [accumulator, *context] -> [accumulator + f(context), *context]
  
  The attention_layer must be an instance of EfficientAttentionBase or one of
  its subclasses (see efficient_attention.py), or None.
  
  Attention is special-cased for the following two reasons:
  
  - LSH attention needs to save bucket assignments from the forward pass to the
    backward pass, for training stability. This requires special-casing it.
  - We can call attention_layer.forward_and_or_backward to compute its output
    (needed for inverting a reversible residual layer) while simultaneously
    performing the backward pass. Sharing computation between these two
    operations improves training speed.
class_docstr_tok: ["Half","of","a","RevNet-style","residual","that","optionally","performs","attention",".","When","attention_layer","is","None",",","this","layer","has","the","signature",":",":","[","accumulator",",","*","context","]","-",">","[","accumulator","+","f","(","context",")",",","*","context","]","The","attention_layer","must","be","an","instance","of","EfficientAttentionBase","or","one","of","its","subclasses","(","see","efficient_attention.py",")",",","or","None",".","Attention","is","special-cased","for","the","following","two","reasons",":","-","LSH","attention","needs","to","save","bucket","assignments","from","the","forward","pass","to","the","backward","pass",",","for","training","stability",".","This","requires","special-casing","it",".","-","We","can","call","attention_layer.forward_and_or_backward","to","compute","its","output","(","needed","for","inverting","a","reversible","residual","layer",")","while","simultaneously","performing","the","backward","pass",".","Sharing","computation","between","these","two","operations","improves","training","speed","."]
code: def __init__(self, *residual_layers, attention_layer=None, name=None):
      super().__init__(name=name)
  
      self._compute_residual = cb.Serial(*residual_layers)
      self._attention_layer = attention_layer
  
      if self._attention_layer is None:
          self._sublayers = (self._compute_residual,)
      else:
          if hasattr(attention_layer, "forward_and_or_backward"):
              self._forward_and_or_backward = (
                  attention_layer.forward_and_or_backward
              )
          else:
              self._forward_and_or_backward = _forward_and_or_backward(
                  attention_layer
              )
          self._sublayers = (
              self._compute_residual,
              self._attention_layer,
          )
  
      running_max = 0
      running_total = 0
      for layer in self._sublayers:
          running_total += layer.n_in
          running_max = max(running_max, running_total)
          running_total -= layer.n_out
      self._n_in = self._n_out = running_max + 1
  
imports: from absl import logging
  import jax
  from trax import fastmath
  from trax.layers import base
  from trax.layers import combinators
label_desc: Use image node_id 1 to create a new ReversibleHalfResidual object from inherited base classes: ReversibleLayer with example: obj = ReversibleHalfResidual()