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 forward(self, xs):
      rngs = _split_rngs(self.rng, len(self.sublayers))
      accumulator, *context = xs
      stack = context = tuple(context)
      new_state = []
      for layer, w, s, rng in zip(
          self.sublayers, self.weights, self.state, rngs
      ):
          inputs = cb.inputs_from_stack(stack, layer.n_in)
          if base.N_WEIGHTS_SHARDS > 1:
              # With sharded weights, make sure we don't keep them concatenated
              # in memory on each device by using remat.
              outputs, s = jax.remat(layer.pure_fn)(inputs, w, s, rng)
          else:
              outputs, s = layer.pure_fn(inputs, w, s, rng)
          stack = cb.outputs_onto_stack(outputs, stack, layer.n_in)
          new_state.append(s)
      residual = stack[0] if isinstance(stack, (tuple, list)) else stack
  
      output = accumulator + residual
      stack = (output,) + context
      self.state = tuple(new_state)
      return stack
  
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 2 for calling the ReversibleHalfResidual obj's underlying member method code with example usage: obj.forward(xs) and returns: stack