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_weights_and_state(self, input_signature):
      stack = input_signature[1:]
      if len(stack) == 1:
          stack = stack[0]
  
      inputs = cb.inputs_from_stack(stack, self._compute_residual.n_in)
      weights, state = self._compute_residual.init(inputs)
      outputs, _ = self._compute_residual._forward_abstract(inputs)
      stack = cb.outputs_onto_stack(
          outputs, stack, self._compute_residual.n_in
      )
  
      if self._attention_layer is None:
          self.state = (state,)
          self.weights = (weights,)
      else:
          inputs = cb.inputs_from_stack(
              stack, self._attention_layer.n_in
          )
          attn_weights, attn_state = self._attention_layer.init(inputs)
          self.state = (state, attn_state)
          self.weights = (weights, attn_weights)
  
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 5 for calling the ReversibleHalfResidual obj's underlying member method code with example usage: obj.init_weights_and_state(input_signature) without return types