src/transformers/models/gpt2/modeling_flax_gpt2.py/FlaxGPT2Attention/setup
class FlaxGPT2Attention:    def setup(self):
        config = self.config
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.embed_dim // self.num_heads

        if self.is_cross_attention:
            self.c_attn = FlaxConv1D(2 * self.embed_dim, dtype=self.dtype)
            self.q_attn = FlaxConv1D(self.embed_dim, dtype=self.dtype)
        else:
            self.c_attn = FlaxConv1D(3 * self.embed_dim, dtype=self.dtype)
        self.c_proj = FlaxConv1D(self.embed_dim, dtype=self.dtype)

        self.resid_dropout = nn.Dropout(rate=config.resid_pdrop)

        if self.causal:
            self.causal_mask = make_causal_mask(
                jnp.ones((1, config.max_position_embeddings), dtype="bool"), dtype="bool"
            )