src/transformers/generation/flax_logits_process.py/FlaxSuppressTokensAtBeginLogitsProcessor/__call__
class FlaxSuppressTokensAtBeginLogitsProcessor:    def __call__(self, input_ids, scores, cur_len: int):
        apply_penalty = 1 - jnp.bool_(cur_len - self.begin_index)

        scores = jnp.where(apply_penalty, scores.at[:, self.begin_suppress_tokens].set(-float("inf")), scores)

        return scores