def forward(self, xs):
    """Executes this layer as part of a forward pass through the model."""
    rng1, rng2 = fastmath.random.split(self.rng, 2)
    l = xs[0].shape[1] if isinstance(xs, tuple) else xs.shape[1]
    if self._std_length is None or l > self._std_length:
      s = self.state[0]
      outputs, s = self._lsha.pure_fn(xs, (), s, rng1, use_cache=True)
      self.state = (s, self.state[1])
    else:
      s = self.state[1]
      w = ((), ())  # std attention is a Serial(Dup, DotProduct), needs 2 ()
      outputs, s = self._stda.pure_fn(xs, w, s, rng2, use_cache=True)
      self.state = (self.state[0], s)
    return outputs
Executes this layer as part of a forward pass through the model.
trax/layers/research/efficient_attention.py