src/transformers/models/big_bird/modeling_flax_big_bird.py/FlaxBigBirdForQuestionAnswering/prepare_question_mask
class FlaxBigBirdForQuestionAnswering:    def prepare_question_mask(q_lengths, maxlen: int):
        # q_lengths -> (bz, 1)
        mask = jnp.arange(0, maxlen)
        mask = jnp.expand_dims(mask, axis=0) < q_lengths
        return mask