src/transformers/models/gpt_neo/modeling_flax_gpt_neo.py/FlaxGPTNeoAttention/setup
class FlaxGPTNeoAttention:    def setup(self):
        attention_type = self.config.attention_layers[self.layer_id]
        self.attention = FlaxGPTNeoSelfAttention(self.config, attention_type, dtype=self.dtype)