src/transformers/generation/flax_logits_process.py/FlaxForcedBOSTokenLogitsProcessor/__call__
class FlaxForcedBOSTokenLogitsProcessor:    def __call__(self, input_ids: jnp.ndarray, scores: jnp.ndarray, cur_len: int) -> jnp.ndarray:
        new_scores = jnp.full(scores.shape, -float("inf"))

        apply_penalty = 1 - jnp.bool_(cur_len - 1)

        scores = jnp.where(apply_penalty, new_scores.at[:, self.bos_token_id].set(0), scores)

        return scores