examples/research_projects/performer/modeling_flax_performer.py/FlaxPerformerIntermediate/__call__
class FlaxPerformerIntermediate:    def __call__(self, hidden_state):
        # TODO: Add ACT2FN reference to change activation function
        dense = nn.Dense(features=self.output_size, name="dense")(hidden_state)
        return ACT2FN[self.hidden_act](dense)