examples/research_projects/performer/modeling_flax_performer.py/FlaxPerformerEmbedding/__call__
class FlaxPerformerEmbedding:    def __call__(self, inputs):
        embedding = self.param("weight", self.emb_init, (self.vocab_size, self.hidden_size))
        return jnp.take(embedding, inputs, axis=0)