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

    See `forward` for a reference implementation of what this layer does. The
    reference implementation is not very efficient, however, and this method
    provides a more performant version.

    Args:
      inputs: inputs to the attention layer tuple (qk, v, mask)
      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
    """
    # TODO(b/148460708): reduce memory usage further
    # TODO(kitaev): there should be a higher-level API (like vmap) that does
    #     batching, instead of needing 3 separate manual implementations here.

    # Notes regarding the implementation:
    # (a) Multiple heads or examples are batched together. There are three
    #     different regimes possible: one head at a time (for long sequences and
    #     expensive attention types), several attention heads at a time (for
    #     long sequences but less-expensive attention types), and several
    #     examples at a time (for large batches of shorter sequences). For the
    #     time being, each of these regimes has its own code.
    # (b) Python loops produce large computation graphs when jitted, so the
    #     default is to use a JAX loop instead.
    # (c) No intermediate quantities are cached for the backward pass. Instead,
    #     the forward pass is re-computed when doing backprop. This approach is
    #     often called "checkpointing" or "rematerialization". When not all
    #     examples or heads fit in memory simultaneously, the implementation
    #     should be [FW-BW-1] and NOT [FW-BW-2], because the latter has worse
    #     memory locality. I don't think JAX autodiff can synthesize [FW-BW-1]
    #     automatically, so the looping for the backward pass is done manually.
    #
    #     [FW-BW-1] for example, head in zip(examples, heads):
    #                 forward(example, head)
    #                 backward(example, head)  # uses intermediates from forward
    #
    #     [FW-BW-2] for example, head in zip(examples, heads):
    #                 forward(example, head)
    #               for example, head in zip(examples, heads):
    #                 backward(example, head)

    if self._masked:
      qk, v, mask = inputs
      batch_size = mask.shape[0]
    else:
      qk, v = inputs
      mask = None
      batch_size = qk.shape[0] // self._n_heads
    batch_x_heads, seqlen, d_model = qk.shape

    compute_grad = (output_grad is not None)
    assert compute_output or compute_grad, 'No work to perform!'

    if not self._incremental:
      forward_unbatched = functools.partial(
          self.forward_unbatched, rng=rng, update_state=update_state)
    else:
      assert not compute_grad

      # The input to use_predict_mem is (qk, v)
      inputs = (qk, v)

      if update_state:
        inputs, state, q_start, new_mem, new_mem_end = self._use_predict_mem(
            inputs, state)
      else:
        # This assumes that the memory stores all of the inputs, which would not
        # be valid if doing backprop in mode 'predict' with long lengths.
        new_mem_end, inputs, state = state
        q_start = new_mem_end - seqlen

      # Reset qk and v to what use_predict_mem/state gave us.
      qk, v = inputs

      forward_unbatched = functools.partial(
          self._incremental_forward_unbatched,
          q_start=fastmath.stop_gradient(q_start),
          q_len=fastmath.stop_gradient(seqlen),
          rng=rng, update_state=update_state)

    # Adjust degree of parallelism based on the batch size.
    n_parallel_heads = batch_size * self._n_heads
    if self._n_parallel_heads and self._n_parallel_heads < n_parallel_heads:
      n_parallel_heads = self._n_parallel_heads

    def tree_update(tree, indices, new_values):
      return fastmath.nested_map_multiarg(
          lambda x, y: fastmath.index_update(x, jax.numpy.index_exp[indices],
                                             y),
          tree, new_values)

    def tree_add(tree, indices, addends):
      return fastmath.nested_map_multiarg(
          lambda x, y: fastmath.index_add(x, jax.numpy.index_exp[indices], y),
          tree, addends)

    if compute_grad:
      inputs_is_differentiable = fastmath.nested_map(
          lambda x: np.issubdtype(x.dtype, np.inexact), inputs)
      def split_differentiable(xs):
        differentiable_xs = fastmath.nested_map_multiarg(
            lambda x, is_differentiable: x if is_differentiable else None,
            xs, inputs_is_differentiable)
        non_differentiable_xs = fastmath.nested_map_multiarg(
            lambda x, is_differentiable: None if is_differentiable else x,
            xs, inputs_is_differentiable)
        return differentiable_xs, non_differentiable_xs
      def join_differentiable(differentiable_xs, non_differentiable_xs):
        """Reconstitute inputs pytree from differentiable/non-d. partitions."""
        differentiable_leaves = fastmath.tree_leaves(differentiable_xs)
        non_differentiable_leaves = fastmath.tree_leaves(non_differentiable_xs)
        leaves = []
        for is_differentiable in fastmath.tree_leaves(inputs_is_differentiable):
          if is_differentiable:
            leaves.append(differentiable_leaves.pop(0))
          else:
            leaves.append(non_differentiable_leaves.pop(0))
        assert not differentiable_leaves
        assert not non_differentiable_leaves
        tree, _ = fastmath.tree_unflatten(leaves, inputs)
        return tree

      def vjp(fn, inp, *args, has_aux=False):
        d_inp, nd_inp = split_differentiable(inp)

        def fn_closed_over_nd_inp(d_inp, *args):
          inp = join_differentiable(d_inp, nd_inp)
          return fn(inp, *args)
        return fastmath.vjp(fn_closed_over_nd_inp, d_inp, *args,
                            has_aux=has_aux)

    if n_parallel_heads != 1:
      raise NotImplementedError(
          'PureLSHSelfAttention is not implemented for n_parallel_heads != 1.')

    def run_inner(idx, loop_val):
      """Runs one slice of attention (for a single head)."""
      o_all, s_all, i_ct_all = loop_val
      example_idx = idx // self._n_heads
      unused_head_idx = idx % self._n_heads

      s_h = fastmath.nested_map(lambda s: s[idx], state)

      if self._masked:
        i_h = (qk[idx], v[idx], mask[example_idx])
      else:
        i_h = (qk[idx], v[idx])

      def forward_fn(i_h):
        return forward_unbatched(
            *i_h, state=fastmath.stop_gradient(s_h))

      if compute_grad:
        o_h, backward_fn, s_h = vjp(forward_fn, i_h, has_aux=True)
        ct_h = output_grad[idx]
        assert o_h.shape == ct_h.shape
        i_ct_h, = backward_fn(ct_h)
      else:
        o_h, s_h = forward_fn(i_h)

      if compute_output:
        o_all = tree_update(o_all, idx, o_h)
      if update_state:
        s_all = tree_update(s_all, idx, s_h)
      if compute_grad:
        i_ct_all = tree_add(i_ct_all, idx, i_ct_h)
      return (o_all, s_all, i_ct_all)

    o_all = s_all = i_ct_all = None
    if compute_output:
      o_all = np.zeros((batch_x_heads, seqlen, d_model), dtype=v.dtype)
    if update_state:
      s_all = state
    if compute_grad:
      i_ct_all = fastmath.nested_map(np.zeros_like, inputs)
      i_ct_all, i_nondifferentiable_dummy_ct = split_differentiable(i_ct_all)

    loop_val = (o_all, s_all, i_ct_all)

    assert (batch_size * self._n_heads) % n_parallel_heads == 0
    loop_hi = (batch_size * self._n_heads) // n_parallel_heads
    if self._use_python_loop or loop_hi == 1:
      for idx in range(loop_hi):
        loop_val = run_inner(idx, loop_val)
    else:
      loop_val = fastmath.fori_loop(
          0, loop_hi, run_inner, loop_val)

    (o_all, s_all, i_ct_all) = loop_val

    if compute_grad:
      i_ct_all = join_differentiable(i_ct_all, i_nondifferentiable_dummy_ct)

    if self._incremental and update_state:
      s_all = (new_mem_end, new_mem, s_all)

    return (o_all, s_all, i_ct_all)
Performs batched forward and/or backward passes.

    See `forward` for a reference implementation of what this layer does. The
    reference implementation is not very efficient, however, and this method
    provides a more performant version.

    Args:
      inputs: inputs to the attention layer tuple (qk, v, mask)
      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