src/transformers/models/vit/modeling_flax_vit.py/FlaxViTAttention/setup
class FlaxViTAttention:    def setup(self):
        self.attention = FlaxViTSelfAttention(self.config, dtype=self.dtype)
        self.output = FlaxViTSelfOutput(self.config, dtype=self.dtype)