examples/research_projects/jax-projects/big_bird/bigbird_flax.py/FlaxBigBirdForNaturalQuestionsModule/__call__
class FlaxBigBirdForNaturalQuestionsModule:    def __call__(self, *args, **kwargs):
        outputs = super().__call__(*args, **kwargs)
        cls_out = self.cls(outputs[2])
        return outputs[:2] + (cls_out,)