src/transformers/models/bert/modeling_flax_bert.py/FlaxBertAttention/setup
class FlaxBertAttention:    def setup(self):
        self.self = FlaxBertSelfAttention(self.config, causal=self.causal, dtype=self.dtype)
        self.output = FlaxBertSelfOutput(self.config, dtype=self.dtype)