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).
    """
    state, rng = self.state, self.rng

    if self._use_reference_code:
      raise NotImplementedError(
          'Reference code not implemented for PureLSHSelfAttention')

    output, new_state, unused_input_cotangents = self.forward_and_or_backward(
        inputs, state, rng, compute_output=True, update_state=True)
    self.state = 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