def forward(self, inputs):
    """Computes this layer's output as part of a forward pass through the model.

    Args:
      inputs: Layer inputs (subclasses may use different inputs)

    Returns:
      A tuple (output, new_state).
    """
    weights, state, rng = self.weights, self.state, self.rng
    if not self._use_reference_code:
      # By default, an efficient, batched implementation is used.
      output, new_state, _, _ = self.forward_and_or_backward(
          inputs, weights, state, rng, compute_output=True, update_state=True)
      self.state = new_state
      return output

    # The reference implementation below provides a more readable overview of
    # what this class does. It's not optimized, however, and should only be used
    # when testing this class for correctness.
    if not isinstance(inputs, (tuple, list)):
      inputs = (inputs,)
    batch_size = int(inputs[0].shape[0])
    seqlen = inputs[0].shape[-2]
    d_model = inputs[0].shape[-1]

    if self._incremental:
      inputs, state, q_start, new_mem, new_mem_end = self._use_predict_mem(
          inputs, state)

    output_accum = [np.zeros((seqlen, d_model)) for _ in range(batch_size)]
    new_state = []
    for example_idx in range(batch_size):
      for head_idx in range(self._n_heads):
        # pylint: disable=cell-var-from-loop
        single_inputs = fastmath.nested_map(lambda x: x[example_idx], inputs)
        single_weights = fastmath.nested_map(lambda w: w[head_idx], weights)
        single_state = fastmath.nested_map(
            lambda s: s[example_idx * self._n_heads + head_idx], state)
        # pylint: enable=cell-var-from-loop
        if self._incremental:
          single_out, single_new_state = self._incremental_forward_unbatched(
              *single_inputs, q_start=q_start, q_len=seqlen,
              weights=single_weights, rng=rng,
              state=single_state, update_state=True)
        else:
          single_out, single_new_state = self.forward_unbatched(
              *single_inputs, weights=single_weights, rng=rng,
              state=single_state, update_state=True)
        new_state.append(single_new_state)
        output_accum[example_idx] = output_accum[example_idx] + single_out

    output = np.stack(output_accum, 0)
    if new_state and fastmath.tree_leaves(new_state[0]):
      new_state = fastmath.nested_map_multiarg(
          lambda *s: np.stack(s, 0), *new_state)
    else:
      new_state = state
    if self._incremental:
      new_state = (new_mem_end, new_mem, new_state)
    self.state = tuple(new_state)
    return output
Computes this layer's output as part of a forward pass through the model.

    Args:
      inputs: Layer inputs (subclasses may use different inputs)

    Returns:
      A tuple (output, new_state).
    
trax/layers/research/efficient_attention.py