def forward_and_or_backward(self, 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
    """
    assert compute_output
    assert not update_state
    assert output_grad is not None

    rngs = fastmath.random.split(rng, 4)
    # Layer order forward: self._qkv, self._attn, self._merge, self._dense
    # Use forward_and_or_backward for attn.

    qkv_output, qkv_vjp_fn, unused_qkv_new_state = fastmath.vjp(
        self._qkv.pure_fn, inputs, weights[0], state[0],
        rngs[0], has_aux=True)

    attn_output, _, _ = self._attn.forward_and_or_backward(
        qkv_output, state[1], rngs[1], output_grad=None,
        compute_output=True, update_state=False)

    merge_output, merge_vjp_fn, unused_merge_new_state = fastmath.vjp(
        self._merge.pure_fn, attn_output, weights[2], state[2], rngs[2],
        has_aux=True)

    if self._dense is not None:
      dense_output, dense_vjp_fn, unused_dense_new_state = fastmath.vjp(
          self._dense.pure_fn, merge_output, weights[3], state[3],
          rngs[3], has_aux=True)

    # Now backward.
    if self._dense is not None:
      dense_grads_inputs, dense_grads_weights, _, _ = dense_vjp_fn(
          output_grad)
    else:
      dense_grads_inputs = output_grad
    merge_grads_inputs, merge_grads_weights, _, _ = merge_vjp_fn(
        dense_grads_inputs)

    # Use forward_and_or_backward for attn.
    (attn_output, _, attn_grads_inputs) = self._attn.forward_and_or_backward(
        qkv_output, state[1], rngs[1], output_grad=merge_grads_inputs,
        compute_output=True, update_state=False)

    # Backward for qkv layer.
    qkv_grad_inputs, qkv_grads_weights, _, _ = qkv_vjp_fn(attn_grads_inputs)

    if self._dense is None:
      grads_weights = (qkv_grads_weights,
                       (),
                       merge_grads_weights)
    else:
      grads_weights = (qkv_grads_weights,
                       (),
                       merge_grads_weights,
                       dense_grads_weights)

    # Output is  (output, new_state,  inputs_grad, weights_grad).
    # new_state is None because update_state is False.
    if self._dense is None:
      return (merge_output, None, qkv_grad_inputs, grads_weights)
    else:
      return (dense_output, None, qkv_grad_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/research/efficient_attention.py