src/transformers/models/mbart/modeling_flax_mbart.py/FlaxMBartForSequenceClassificationModule/__call__
class FlaxMBartForSequenceClassificationModule:    def __call__(
        self,
        input_ids,
        attention_mask,
        decoder_input_ids,
        decoder_attention_mask,
        position_ids,
        decoder_position_ids,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
        deterministic: bool = True,
    ):
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            decoder_input_ids=decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask,
            position_ids=position_ids,
            decoder_position_ids=decoder_position_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=deterministic,
        )

        hidden_states = outputs[0]  # last hidden state

        eos_mask = jnp.where(input_ids == self.config.eos_token_id, 1, 0)

        # The first condition is necessary to overcome jax._src.errors.ConcretizationTypeError during JIT compilation
        if type(eos_mask) != jax.interpreters.partial_eval.DynamicJaxprTracer:
            if len(jnp.unique(eos_mask.sum(1))) > 1:
                raise ValueError("All examples must have the same number of <eos> tokens.")

            if any(eos_mask.sum(1) == 0):
                raise ValueError("There are missing <eos> tokens in input_ids")

            # Ensure to keep 1 only for the last <eos> token for each example
            eos_mask_noised = eos_mask + jnp.arange(eos_mask.shape[1]) * 1e-6
            eos_mask = jnp.where(eos_mask_noised == eos_mask_noised.max(1).reshape(-1, 1), 1, 0)

        sentence_representation = jnp.einsum("ijk, ij -> ijk", hidden_states, eos_mask).sum(1)
        logits = self.classification_head(sentence_representation, deterministic=deterministic)

        if not return_dict:
            output = (logits,) + outputs[1:]
            return output

        return FlaxSeq2SeqSequenceClassifierOutput(
            logits=logits,
            decoder_hidden_states=outputs.decoder_hidden_states,
            decoder_attentions=outputs.decoder_attentions,
            cross_attentions=outputs.cross_attentions,
            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
            encoder_hidden_states=outputs.encoder_hidden_states,
            encoder_attentions=outputs.encoder_attentions,
        )