class_docstr: Reversible Layer.
class_docstr_tok: ["Reversible","Layer","."]
docstr: Backward pass: computes the inverse of a layer and propagates gradients.
  
  While you may choose to only implement reverse, some layers implement this
  function directly as computation may be shared between reversing and
  computing gradients.
  
  Args:
    output: Output activations; can be a (possibly nested) tuple.
    grad: gradient signal (cotangent) computed based on subsequent layers.
      The structure and shape must match the output.
    weights: layer weights
    state: start state
    new_state: updated state computed by the forward pass
    rng: Single-use random number generator (JAX PRNG key).
  
  Returns:
    A tuple (x, (x_grad, weights_grad)), where x is the reconstructed input,
    x_grad is the gradient signal for the input, and weights_grad is the
    gradient signal for the weights.
docstr_tok: ["Backward","pass",":","computes","the","inverse","of","a","layer","and","propagates","gradients",".","While","you","may","choose","to","only","implement","reverse",",","some","layers","implement","this","function","directly","as","computation","may","be","shared","between","reversing","and","computing","gradients",".","Args",":","output",":","Output","activations",";","can","be","a","(","possibly","nested",")","tuple",".","grad",":","gradient","signal","(","cotangent",")","computed","based","on","subsequent","layers",".","The","structure","and","shape","must","match","the","output",".","weights",":","layer","weights","state",":","start","state","new_state",":","updated","state","computed","by","the","forward","pass","rng",":","Single-use","random","number","generator","(","JAX","PRNG","key",")",".","Returns",":","A","tuple","(","x",",","(","x_grad",",","weights_grad",")",")",",","where","x","is","the","reconstructed","input",",","x_grad","is","the","gradient","signal","for","the","input",",","and","weights_grad","is","the","gradient","signal","for","the","weights","."]
code: def reverse_and_grad(
      self, output, grad, weights=(), state=(), new_state=(), rng=None
  ):
      """Backward pass: computes the inverse of a layer and propagates gradients.
  
      While you may choose to only implement reverse, some layers implement this
      function directly as computation may be shared between reversing and
      computing gradients.
  
      Args:
        output: Output activations; can be a (possibly nested) tuple.
        grad: gradient signal (cotangent) computed based on subsequent layers.
          The structure and shape must match the output.
        weights: layer weights
        state: start state
        new_state: updated state computed by the forward pass
        rng: Single-use random number generator (JAX PRNG key).
  
      Returns:
        A tuple (x, (x_grad, weights_grad)), where x is the reconstructed input,
        x_grad is the gradient signal for the input, and weights_grad is the
        gradient signal for the weights.
      """
      reconstructed_x = self.reverse(
          output, weights, state, new_state, rng
      )
      _, vjpfun = fastmath.vjp(
          self._pure_forward, reconstructed_x, weights, state, rng
      )
      x_grad, weights_grad, _, _ = vjpfun(grad)
      return reconstructed_x, (x_grad, 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 ReversibleLayer obj's underlying member method code with example usage: obj.reverse_and_grad(output, grad, weights, state, new_state, rng) and returns: reconstructed_x,  x_grad,  weights_grad