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.

    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
      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
    """
    # TODO(kitaev): profile ~4% speed drop compared to previous implementation
    #     in some conditions. Other conditions (e.g. the enwik8 model) appear
    #     to have the same overall training speed.
    # 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)

    have_single_input = not isinstance(inputs, (tuple, list))
    if have_single_input:
      inputs = (inputs,)
    batch_size = int(inputs[0].shape[0])
    seqlen = inputs[0].shape[-2]
    d_model = inputs[0].shape[-1]

    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:
      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

      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, new_values):
      return fastmath.nested_map_multiarg(
          lambda x, y: fastmath.index_add(x, jax.numpy.index_exp[indices], y),
          tree, new_values)

    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:
      def run_inner(idx, loop_val):
        """Runs one slice of attention (for a single head)."""
        o_all, s_all, i_ct_all, w_ct_all = loop_val
        example_idx = idx // self._n_heads
        head_idx = idx % self._n_heads

        i_h = fastmath.nested_map(lambda x: x[example_idx], inputs)
        w_h = fastmath.nested_map(lambda w: w[head_idx], weights)
        s_h = fastmath.nested_map(lambda s: s[idx], state)

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

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

        if compute_output:
          o_all = fastmath.index_add(o_all, example_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, example_idx, i_ct_h)
          w_ct_all = tree_add(w_ct_all, head_idx, w_ct_h)
        return (o_all, s_all, i_ct_all, w_ct_all)
    elif n_parallel_heads < self._n_heads:
      assert self._n_heads % n_parallel_heads == 0
      def run_inner(idx, loop_val):
        """Runs one slice of attention (multiple heads, but one example)."""
        o_all, s_all, i_ct_all, w_ct_all = loop_val
        idx = idx * self._n_parallel_heads
        example_idx = idx // self._n_heads
        head_idx_lo = idx % self._n_heads
        head_range = head_idx_lo + np.arange(n_parallel_heads, dtype=np.int32)
        state_range = idx + np.arange(n_parallel_heads, dtype=np.int32)

        i_mh = fastmath.nested_map(lambda x: x[example_idx], inputs)
        w_mh = fastmath.nested_map(lambda w: w[head_range], weights)
        s_mh = fastmath.nested_map(lambda s: s[state_range], state)
        def forward_unbatched_h(i_h, w_h, s_h):
          return forward_unbatched(*i_h, weights=w_h, state=s_h)
        def forward_fn(i_mh, w_mh):
          o_mh, new_s_mh = fastmath.vmap(
              forward_unbatched_h, in_axes=(None, 0, 0), out_axes=0)(
                  i_mh, w_mh, s_mh)
          o_mh = np.sum(o_mh, axis=0)
          return o_mh, new_s_mh

        if compute_grad:
          o_mh, backward_fn, s_mh = vjp(forward_fn, i_mh, w_mh, has_aux=True)
          ct_mh = output_grad[example_idx]
          assert o_mh.shape == ct_mh.shape
          i_ct_mh, w_ct_mh = backward_fn(ct_mh)
        else:
          o_mh, s_mh = forward_fn(i_mh, w_mh)

        if compute_output:
          o_all = fastmath.index_add(o_all, example_idx, o_mh)
        if update_state:
          s_all = tree_update(s_all, state_range, s_mh)
        if compute_grad:
          i_ct_all = tree_add(i_ct_all, example_idx, i_ct_mh)
          w_ct_all = tree_add(w_ct_all, head_range, w_ct_mh)
        return (o_all, s_all, i_ct_all, w_ct_all)
    else:
      assert n_parallel_heads % self._n_heads == 0
      def forward_single_example(i_x, w_all, s_x):
        def forward_unbatched_h(i_h, w_h, s_h):
          return forward_unbatched(*i_h, weights=w_h, state=s_h)
        o_x, s_x = fastmath.vmap(
            forward_unbatched_h, in_axes=(None, 0, 0), out_axes=(0, 0))(
                i_x, w_all, s_x)
        o_x = np.sum(o_x, axis=0)
        return o_x, s_x
      def run_inner(idx, loop_val):
        """Runs one slice of attention (all heads for one or more examples)."""
        o_all, s_all, i_ct_all, w_ct_all = loop_val
        idx = idx * n_parallel_heads
        example_idx_lo = idx // self._n_heads
        example_range = example_idx_lo + np.arange(
            n_parallel_heads // self._n_heads, dtype=np.int32)
        state_range = idx + np.arange(n_parallel_heads, dtype=np.int32)

        i_mex = fastmath.nested_map(lambda x: x[example_range], inputs)
        s_mex = fastmath.nested_map(
            lambda s: np.reshape(s[state_range],  # pylint: disable=g-long-lambda
                                 (-1, self._n_heads) + s.shape[1:]),
            state)
        def forward_fn(i_mex, w_all):
          o_mex, new_s_mex = fastmath.vmap(
              forward_single_example, in_axes=(0, None, 0), out_axes=(0, 0))(
                  i_mex, w_all, s_mex)
          new_s_mex = fastmath.nested_map(
              lambda s: np.reshape(s, (n_parallel_heads,) + s.shape[2:]),
              new_s_mex)
          return o_mex.astype(i_mex[0].dtype), new_s_mex

        if compute_grad:
          o_mex, backward_fn, s_mex = vjp(forward_fn, i_mex, weights,
                                          has_aux=True)
          ct_mex = output_grad[example_range]
          assert o_mex.shape == ct_mex.shape, str(ct_mex.shape)
          assert o_mex.dtype == ct_mex.dtype, str(ct_mex.dtype)
          i_ct_mex, w_ct_mex = backward_fn(ct_mex)
        else:
          o_mex, s_mex = forward_fn(i_mex, weights)

        if compute_output:
          o_all = fastmath.index_add(o_all, jax.numpy.index_exp[example_range],
                                     o_mex)
        if update_state:
          s_all = tree_update(s_all, state_range, s_mex)
        if compute_grad:
          i_ct_all = tree_update(i_ct_all, example_range, i_ct_mex)
          w_ct_all = fastmath.nested_map_multiarg(
              lambda old_all, delta_all: old_all + delta_all,
              w_ct_all, w_ct_mex)
        return (o_all, s_all, i_ct_all, w_ct_all)

    o_all = s_all = i_ct_all = w_ct_all = None
    if compute_output:
      o_all = np.zeros(
          (batch_size, seqlen, d_model), dtype=inputs[0].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)
      w_ct_all = fastmath.nested_map(np.zeros_like, weights)

    loop_val = (o_all, s_all, i_ct_all, w_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, w_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)

    if have_single_input and compute_grad:
      assert isinstance(i_ct_all, tuple) and len(i_ct_all) == 1
      return (o_all, s_all, i_ct_all[0], w_ct_all)
    else:
      return (o_all, s_all, i_ct_all, w_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
      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