src/transformers/models/marian/modeling_flax_marian.py/FlaxMarianMTModel/_adapt_logits_for_beam_search
class FlaxMarianMTModel:    def _adapt_logits_for_beam_search(self, logits):
        """This function enforces the padding token never to be generated."""
        logits = jax.ops.index_update(logits, jax.ops.index[:, :, self.config.pad_token_id], float("-inf"))
        return logits