class_docstr: A reversible version of tl.Serial (requires reversible sub-layers).
class_docstr_tok: ["A","reversible","version","of","tl.Serial","(","requires","reversible","sub-layers",")","."]
code: def reverse_and_grad(
      self, output, grad, weights=(), state=(), new_state=(), rng=None
  ):
      rngs = (None,) * self._n_layers
      if rng is not None:
          rngs = fastmath.random.split(rng, self._n_layers)
  
      stack = output
      stack_grad = grad
      weights_grad = []
      for layer, p, s, ns, rng in reversed(
          list(zip(self.sublayers, weights, state, new_state, rngs))
      ):
          layer_val = cb.inputs_from_stack(stack, layer.n_out)
          layer_ct = cb.inputs_from_stack(stack_grad, layer.n_out)
          layer_val, layer_ct = layer.reverse_and_grad(
              layer_val, layer_ct, p, s, ns, rng=rng
          )
          layer_ct, p_ct = layer_ct
          weights_grad.insert(0, p_ct)
          stack = cb.outputs_onto_stack(layer_val, stack, layer.n_out)
          stack_grad = cb.outputs_onto_stack(
              layer_ct, stack_grad, layer.n_out
          )
  
      return stack, (stack_grad, tuple(weights_grad))
  
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 3 for calling the ReversibleSerial obj's underlying member method code with example usage: obj.reverse_and_grad(output, grad, weights, state, new_state, rng) and returns: stack,  stack_grad,  tuple