examples/research_projects/performer/modeling_flax_performer.py/FlaxPerformerLayer/__call__
class FlaxPerformerLayer:    def __call__(self, hidden_state, attention_mask):
        attention = FlaxPerformerAttention(self.num_heads, self.head_size, name="attention")(
            hidden_state, attention_mask
        )
        intermediate = FlaxPerformerIntermediate(
            self.intermediate_size, name="intermediate", hidden_act=self.hidden_act
        )(attention)
        output = FlaxPerformerOutput(name="output")(intermediate, attention)

        return output