examples/research_projects/jax-projects/big_bird/bigbird_flax.py/FlaxBigBirdForNaturalQuestionsModule/setup
class FlaxBigBirdForNaturalQuestionsModule:    def setup(self):
        super().setup()
        self.cls = nn.Dense(5, dtype=self.dtype)