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."""
    assert compute_output
    assert not update_state

    l = inputs[0].shape[1] if isinstance(inputs, tuple) else inputs.shape[1]
    rng1, rng2 = fastmath.random.split(rng, 2)
    if self._std_length is None or l > self._std_length:
      # Run the LSH layer
      s = state[0]
      (out, unused_new_s, grads_inputs) = self._lsha.forward_and_or_backward(
          inputs, s, rng1, output_grad=output_grad,
          compute_output=True, update_state=False)
    else:
      # Run the standard layer
      s, w = state[1], ((), ())
      out, std_vjp_fn, unused_new_s = fastmath.vjp(
          self._stda.pure_fn, inputs, w, s, rng2, has_aux=True)
      if output_grad is not None:
        grads_inputs, _, _, _ = std_vjp_fn(output_grad)
      else:
        grads_inputs = None

    return (out, None, grads_inputs)
Performs batched forward and/or backward passes.
trax/layers/research/efficient_attention.py