src/transformers/generation_flax_logits_process.py/FlaxForcedEOSTokenLogitsProcessor/__call__
class FlaxForcedEOSTokenLogitsProcessor:    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 - self.max_length + 1)

        scores = jnp.where(
            apply_penalty, jax.ops.index_update(new_scores, jax.ops.index[:, self.eos_token_id], 0), scores
        )

        return scores