def forward_and_or_backward(inputs, weights, state, rng, output_grad=None,
                              compute_output=True, update_state=True):
    """Performs batched forward and/or backward passes.

    Args:
      inputs: inputs to the attention layer
      weights: weights for the attention layer
      state: state of the attention layer
      rng: PRNG key for the layer (shared across all examples and heads)
      output_grad: gradient of the loss wrt the output of the layer, or None.
          This function performs the backward pass iff `output_grad` is not
          None.
      compute_output: bool: whether to return the output of the forward pass
          (for example, a pure backwards pass does not need to return the
          output).
      update_state: bool: whether to return an updated layer state.

    Returns:
      A tuple (output, new_state, inputs_grad, weights_grad).
      - output is not None iff compute_output is True
      - new_state is not None iff update_state is True
      - inputs_grad & weights_grad are not None iff output_grad is not None
    """
    # Calculate the vector-Jacobian product of the layer pure_fn.
    output, vjp_fn, new_state = fastmath.vjp(
        layer.pure_fn, inputs, weights, state, rng, has_aux=True)
    output = output if compute_output else None
    new_state = new_state if update_state else None

    # The vjp function returns gradients with respect to inputs and weights.
    if output_grad is not None:
      grads_inputs, grads_weights, _, _ = vjp_fn(output_grad)
    else:
      grads_inputs, grads_weights = None, None

    return (output, new_state, grads_inputs, grads_weights)
Performs batched forward and/or backward passes.

    Args:
      inputs: inputs to the attention layer
      weights: weights for the attention layer
      state: state of the attention layer
      rng: PRNG key for the layer (shared across all examples and heads)
      output_grad: gradient of the loss wrt the output of the layer, or None.
          This function performs the backward pass iff `output_grad` is not
          None.
      compute_output: bool: whether to return the output of the forward pass
          (for example, a pure backwards pass does not need to return the
          output).
      update_state: bool: whether to return an updated layer state.

    Returns:
      A tuple (output, new_state, inputs_grad, weights_grad).
      - output is not None iff compute_output is True
      - new_state is not None iff update_state is True
      - inputs_grad & weights_grad are not None iff output_grad is not None
    
trax/layers/reversible.py