def __call__(
        self,
        input_ids: jnp.ndarray,
        attention_mask: Optional[jnp.ndarray] = None,
        decoder_input_ids: Optional[jnp.ndarray] = None,
        decoder_attention_mask: Optional[jnp.ndarray] = None,
        position_ids: Optional[jnp.ndarray] = None,
        decoder_position_ids: Optional[jnp.ndarray] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        train: bool = False,
        params: dict = None,
        dropout_rng: PRNGKey = None,
    ):
        r"""
        Returns:

        Examples:

        ```python
        >>> from transformers import FlaxEncoderDecoderModel, BertTokenizer, GPT2Tokenizer

        >>> # load a fine-tuned bert2gpt2 model
        >>> model = FlaxEncoderDecoderModel.from_pretrained("patrickvonplaten/bert2gpt2-cnn_dailymail-fp16")
        >>> # load input & output tokenizer
        >>> tokenizer_input = BertTokenizer.from_pretrained("bert-base-cased")
        >>> tokenizer_output = GPT2Tokenizer.from_pretrained("gpt2")

        >>> article = '''Sigma Alpha Epsilon is under fire for a video showing party-bound fraternity members
        >>> singing a racist chant. SAE's national chapter suspended the students,
        >>> but University of Oklahoma President David Boren took it a step further,
        >>> saying the university's affiliation with the fraternity is permanently done.'''

        >>> input_ids = tokenizer_input(article, add_special_tokens=True, return_tensors="np").input_ids

        >>> # use GPT2's eos_token as the pad as well as eos token
        >>> model.config.eos_token_id = model.config.decoder.eos_token_id
        >>> model.config.pad_token_id = model.config.eos_token_id

        >>> sequences = model.generate(input_ids, num_beams=4, max_length=12).sequences

        >>> summary = tokenizer_output.batch_decode(sequences, skip_special_tokens=True)[0]
        >>> assert summary == "SAS Alpha Epsilon suspended Sigma Alpha Epsilon members"
        ```
        """

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.return_dict

        # prepare encoder inputs
        if attention_mask is None:
            attention_mask = jnp.ones_like(input_ids)
        if position_ids is None:
            batch_size, sequence_length = input_ids.shape
            position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))

        # prepare decoder inputs
        if decoder_input_ids is None:
            raise ValueError(
                "`decoder_input_ids` cannot be `None`. For sequence to sequence training, `decoder_position_ids` must"
                " be specified as an input argument."
            )
        if decoder_attention_mask is None:
            decoder_attention_mask = jnp.ones_like(decoder_input_ids)
        if decoder_position_ids is None:
            batch_size, sequence_length = decoder_input_ids.shape
            decoder_position_ids = jnp.broadcast_to(
                jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
            )

        # Handle any PRNG if needed
        rngs = {"dropout": dropout_rng} if dropout_rng is not None else {}

        return self.module.apply(
            {"params": params or self.params},
            input_ids=jnp.array(input_ids, dtype="i4"),
            attention_mask=jnp.array(attention_mask, dtype="i4"),
            decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
            decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
            position_ids=jnp.array(position_ids, dtype="i4"),
            decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            deterministic=not train,
            rngs=rngs,
        )

        Returns:

        Examples:

        ```python
        >>> from transformers import FlaxEncoderDecoderModel, BertTokenizer, GPT2Tokenizer

        >>> # load a fine-tuned bert2gpt2 model
        >>> model = FlaxEncoderDecoderModel.from_pretrained("patrickvonplaten/bert2gpt2-cnn_dailymail-fp16")
        >>> # load input & output tokenizer
        >>> tokenizer_input = BertTokenizer.from_pretrained("bert-base-cased")
        >>> tokenizer_output = GPT2Tokenizer.from_pretrained("gpt2")

        >>> article = '''Sigma Alpha Epsilon is under fire for a video showing party-bound fraternity members
        >>> singing a racist chant. SAE's national chapter suspended the students,
        >>> but University of Oklahoma President David Boren took it a step further,
        >>> saying the university's affiliation with the fraternity is permanently done.'''

        >>> input_ids = tokenizer_input(article, add_special_tokens=True, return_tensors="np").input_ids

        >>> # use GPT2's eos_token as the pad as well as eos token
        >>> model.config.eos_token_id = model.config.decoder.eos_token_id
        >>> model.config.pad_token_id = model.config.eos_token_id

        >>> sequences = model.generate(input_ids, num_beams=4, max_length=12).sequences

        >>> summary = tokenizer_output.batch_decode(sequences, skip_special_tokens=True)[0]
        >>> assert summary == "SAS Alpha Epsilon suspended Sigma Alpha Epsilon members"
        ```
        
dpkd/transformers/src/transformers/models/encoder_decoder/modeling_flax_encoder_decoder.py