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 reverse_and_grad(
      self, output, ct, weights=(), state=(), new_state=(), rng=None
  ):
      rngs = _split_rngs(rng, len(self.sublayers))
  
      accumulator_output, *context = output
      context = tuple(context)
      accumulator_output_ct, *context_ct = ct
      context_ct = tuple(context_ct)
  
      # Forward pass through self._compute_residual. Outputs that will not receive
      # a gradient signal from subsequent layers are moved to aux.
      def call_compute_residual(x, weights):
          state_to_pass = state[0]  # old_state
  
          # _replace_second_time is currently used exclusively in _RememberInReverse
          # layer to combat numerical instability in Terraformer when quantizing
          # the mask in SparseFF.
          def _replace_second_time(stt, nstt):
              if (
                  isinstance(stt, tuple)
                  and len(stt) == 2
                  and isinstance(stt[1], dict)
                  and "running_second_time" in stt[1]
              ):
                  return (nstt[0], {"running_second_time_yes": ()})
              elif isinstance(stt, (tuple, list)):
                  assert isinstance(nstt, (tuple, list)) and len(
                      nstt
                  ) == len(stt)
                  return type(stt)(
                      [
                          _replace_second_time(s, ns)
                          for s, ns in zip(stt, nstt)
                      ]
                  )
              else:
                  return stt
  
          state_to_pass = _replace_second_time(
              state_to_pass, new_state[0]
          )
          res, _ = self._compute_residual.pure_fn(
              x, weights=weights, state=state_to_pass, rng=rngs[0]
          )
          if not isinstance(res, (tuple, list)):
              return res, None
          else:
              n_differentiable = 1
              if self._attention_layer is not None:
                  n_differentiable = min(
                      len(res), self._attention_layer.n_in
                  )
              return res[:n_differentiable], res[n_differentiable:]
  
      stack = context
      inputs = cb.inputs_from_stack(stack, self._compute_residual.n_in)
      outputs, compute_residual_vjpfun, outputs_aux = fastmath.vjp(
          call_compute_residual, inputs, weights[0], has_aux=True
      )
      if outputs_aux is not None:
          n_differentiable_outputs = len(outputs)
          outputs = outputs + outputs_aux
      stack = cb.outputs_onto_stack(
          outputs, stack, self._compute_residual.n_in
      )
  
      stack_ct = accumulator_output_ct
      if self._attention_layer is None:
          residual = (
              stack[0] if isinstance(stack, (tuple, list)) else stack
          )
      else:
          inputs = cb.inputs_from_stack(
              stack, self._attention_layer.n_in
          )
          (
              residual,
              _,
              attn_inputs_ct,
              attn_weights_ct,
          ) = self._forward_and_or_backward(
              inputs,
              weights[1],
              new_state[1],
              rngs[1],
              output_grad=accumulator_output_ct,
              compute_output=True,
              update_state=False,
          )
          stack_ct = cb.outputs_onto_stack(
              attn_inputs_ct, stack_ct, self._attention_layer.n_out
          )
  
      compute_residual_ct = cb.inputs_from_stack(
          stack_ct, self._compute_residual.n_out
      )
      if outputs_aux is not None:
          if not isinstance(compute_residual_ct, (tuple, list)):
              compute_residual_ct = (compute_residual_ct,)
          compute_residual_ct = compute_residual_ct[
              :n_differentiable_outputs
          ]
          assert len(compute_residual_ct) == n_differentiable_outputs
      (
          compute_residual_inputs_ct,
          compute_residual_weights_ct,
      ) = compute_residual_vjpfun(compute_residual_ct)
      stack_ct = cb.outputs_onto_stack(
          compute_residual_inputs_ct,
          stack_ct,
          self._compute_residual.n_out,
      )
      if not isinstance(stack_ct, (tuple, list)):
          stack_ct = (stack_ct,)
  
      def _add(x, y):
          # `None` is for TFNP backend, which uses `None` as the gradient of
          # int/bool instead of an array of dtype `float0`.
          if x is None or x.dtype == jax.float0:
              return y
          if y is None or y.dtype == jax.float0:
              return x
          return x + y
  
      stack_ct = (
          (accumulator_output_ct,)
          + fastmath.nested_map_multiarg(
              _add, context_ct[: len(stack_ct)], stack_ct
          )
          + context_ct[len(stack_ct) :]
      )
  
      reconstructed_x = accumulator_output - residual
      stack = (reconstructed_x,) + context
      if self._attention_layer is None:
          weights_ct = (compute_residual_weights_ct,)
      else:
          weights_ct = (compute_residual_weights_ct, attn_weights_ct)
      return stack, (stack_ct, weights_ct)
  
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 4 for calling the ReversibleHalfResidual obj's underlying member method code with example usage: obj.reverse_and_grad(output, ct, weights, state, new_state, rng) and returns: stack,  stack_ct,  weights_ct, unknown, res,  None, res,  res, y, x, nstt,  dict, type, stt