Clean Flax ViT + GPT2-LM script

vit_gpt2/modeling_flax_vit_gpt2_lm.py

@@ -6,39 +6,27 @@ import jax.numpy as jnp
 from flax.core.frozen_dict import FrozenDict, unfreeze
 from jax import lax
 from jax.random import PRNGKey
-from transformers import GPT2Config, FlaxViTModel, ViTConfig
 from transformers.modeling_flax_outputs import (
     FlaxCausalLMOutputWithCrossAttentions,
     FlaxSeq2SeqLMOutput,
     FlaxSeq2SeqModelOutput,
 )
-from transformers.models.bart.modeling_flax_bart import (
-    shift_tokens_right,
-)
+from .configuration_vit_gpt2 import ViTGPT2Config
+from transformers import ViTConfig, GPT2Config
+### TODO: check FlaxPreTrainedModel
+from transformers import FlaxPreTrainedModel, FlaxViTModel
+from transformers.models.vit.modeling_flax_vit import FlaxViTModule
 from .modeling_flax_gpt2 import (
+    FlaxGPT2PreTrainedModel,
     FlaxGPT2Module,
     FlaxGPT2Model,
     FlaxGPT2LMHeadModule,
     FlaxGPT2LMHeadModel,
-    FlaxPreTrainedModel
 )
-from transformers.models.vit.modeling_flax_vit import FlaxViTModule
-
-from .configuration_vit_gpt2 import ViTGPT2Config
 
 
-def shift_tokens_right(input_ids: jnp.ndarray, pad_token_id: int, decoder_start_token_id: int) -> jnp.ndarray:
-    """
-    Shift input ids one token to the right.
-    """
-    shifted_input_ids = jnp.roll(input_ids, 1, axis=-1)
-    shifted_input_ids = jax.ops.index_update(shifted_input_ids, (..., 0), decoder_start_token_id)
-    # replace possible -100 values in labels by `pad_token_id`
-    shifted_input_ids = jnp.where(shifted_input_ids == -100, pad_token_id, shifted_input_ids)
-
-    return shifted_input_ids
-
 class FlaxViTGPT2LMModule(nn.Module):
+    """Play the same role as ``FlaxBartModule`` but with the decoder equipped with a LM head."""
     config: ViTGPT2Config
     dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
@@ -54,16 +42,16 @@ class FlaxViTGPT2LMModule(nn.Module):
         return self.decoder
 
     def __call__(
-            self,
-            pixel_values,
-            input_ids,
-            attention_mask,
-            position_ids,
-            encoder_attention_mask: Optional[jnp.ndarray] = None,
-            output_attentions: bool = False,
-            output_hidden_states: bool = False,
-            return_dict: bool = True,
-            deterministic: bool = True,
+        self,
+        pixel_values,
+        attention_mask,
+        decoder_input_ids,
+        decoder_attention_mask,
+        decoder_position_ids,
+        output_attentions: bool = False,
+        output_hidden_states: bool = False,
+        return_dict: bool = True,
+        deterministic: bool = True,
     ):
         encoder_outputs = self.encoder(
             pixel_values=pixel_values,
@@ -74,11 +62,11 @@ class FlaxViTGPT2LMModule(nn.Module):
         )
 
         decoder_outputs = self.decoder(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            position_ids=decoder_position_ids,
             encoder_hidden_states=encoder_outputs[0],
-            encoder_attention_mask=encoder_attention_mask,
+            encoder_attention_mask=attention_mask,
             deterministic=deterministic,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
@@ -98,10 +86,14 @@ class FlaxViTGPT2LMModule(nn.Module):
             encoder_attentions=encoder_outputs.attentions,
         )
 
+
 class FlaxViTGPT2LMForConditionalGenerationModule(nn.Module):
+    """Play the same role as ``FlaxBartForConditionalGenerationModule`` but with the decoder equipped with a LM head.
+
+       Actually, it is identical to ``FlaxBartForConditionalGenerationModule`` with a different name.
+    """
     config: ViTGPT2Config
     dtype: jnp.dtype = jnp.float32
-    bias_init: Callable[..., jnp.ndarray] = jax.nn.initializers.zeros
 
     def setup(self):
         self.model = FlaxViTGPT2LMModule(config=self.config, dtype=self.dtype)
@@ -115,10 +107,10 @@ class FlaxViTGPT2LMForConditionalGenerationModule(nn.Module):
     def __call__(
         self,
         pixel_values,
-        input_ids,
         attention_mask,
-        position_ids,
-        encoder_attention_mask: Optional[jnp.ndarray] = None,
+        decoder_input_ids,
+        decoder_attention_mask,
+        decoder_position_ids,
         output_attentions: bool = False,
         output_hidden_states: bool = False,
         return_dict: bool = True,
@@ -126,10 +118,10 @@ class FlaxViTGPT2LMForConditionalGenerationModule(nn.Module):
     ):
         outputs = self.model(
             pixel_values=pixel_values,
-            input_ids=input_ids,
             attention_mask=attention_mask,
-            position_ids=position_ids,
-            encoder_attention_mask=encoder_attention_mask,
+            decoder_input_ids=decoder_input_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            decoder_position_ids=decoder_position_ids,
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
@@ -140,6 +132,7 @@ class FlaxViTGPT2LMForConditionalGenerationModule(nn.Module):
 
 
 class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
+    """Play the same role as ``FlaxBartPretrainedModel``"""
     config_class = ViTGPT2Config
     base_model_prefix: str = "model"
     module_class: nn.Module = None
@@ -159,23 +152,23 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
             )
 
         module = self.module_class(config=config, dtype=dtype, **kwargs)
-        super().__init__(
-            config, module, input_shape=input_shape, seed=seed, dtype=dtype
-        )
+        # This will use ``self.init_weights``.
+        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
 
     def init_weights(self, rng: jax.random.PRNGKey, input_shape: Tuple) -> FrozenDict:
-        # init input tensors
-        pixel_values = jax.random.normal(rng, input_shape[0])
-        # # make sure initialization pass will work for FlaxBartForSequenceClassificationModule
-        # input_ids = jax.ops.index_update(input_ids, (..., -1), self.config.eos_token_id)
 
-        input_ids = jnp.zeros(input_shape[1], dtype="i4")
-        attention_mask = jnp.ones_like(input_ids)
+        encoder_input_shape, decoder_input_shape = input_shape
 
-        batch_size, sequence_length = input_ids.shape
-        position_ids = jnp.broadcast_to(
-            jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
-        )
+        # init input tensors
+        pixel_values = jax.random.normal(rng, encoder_input_shape)
+        attention_mask = None
+        decoder_input_ids = jnp.zeros(decoder_input_shape, dtype="i4")
+        # make sure initialization pass will work for FlaxBartForSequenceClassificationModule
+        decoder_input_ids = jax.ops.index_update(decoder_input_ids, (..., -1), self.config.eos_token_id)
+        decoder_attention_mask = jnp.ones_like(decoder_input_ids)
+
+        batch_size, sequence_length = decoder_input_ids.shape
+        decoder_position_ids = jnp.broadcast_to(jnp.arange(sequence_length)[None, :], (batch_size, sequence_length))
 
         params_rng, dropout_rng = jax.random.split(rng)
         rngs = {"params": params_rng, "dropout": dropout_rng}
@@ -183,40 +176,34 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
         return self.module.init(
             rngs,
             pixel_values,
-            input_ids,
             attention_mask,
-            position_ids,
+            decoder_input_ids,
+            decoder_attention_mask,
+            decoder_position_ids,
         )["params"]
 
     def init_cache(self, batch_size, max_length, encoder_outputs):
-
-        input_ids = jnp.ones((batch_size, max_length), dtype="i4")
-        attention_mask = jnp.ones_like(input_ids)
-        position_ids = jnp.broadcast_to(
-            jnp.arange(jnp.atleast_2d(input_ids).shape[-1]),
-            input_ids.shape,
+        # init input variables to retrieve cache
+        decoder_input_ids = jnp.ones((batch_size, max_length), dtype="i4")
+        decoder_attention_mask = jnp.ones_like(decoder_input_ids)
+        decoder_position_ids = jnp.broadcast_to(
+            jnp.arange(jnp.atleast_2d(decoder_input_ids).shape[-1]), decoder_input_ids.shape,
         )
 
-        def _decoder_forward(
-            module,
-            input_ids,
-            attention_mask,
-            position_ids,
-            **kwargs,
-        ):
+        def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
             decoder_module = module._get_decoder_module()
             return decoder_module(
-                input_ids,
-                attention_mask,
-                position_ids,
+                input_ids=decoder_input_ids,
+                attention_mask=decoder_attention_mask,
+                position_ids=decoder_position_ids,
                 **kwargs,
             )
 
         init_variables = self.module.init(
             jax.random.PRNGKey(0),
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
+            decoder_input_ids=decoder_input_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            decoder_position_ids=decoder_position_ids,
             encoder_hidden_states=encoder_outputs[0],
             init_cache=True,
             method=_decoder_forward,  # we only need to call the decoder to init the cache
@@ -234,20 +221,13 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
         params: dict = None,
         dropout_rng: PRNGKey = None,
     ):
-        output_attentions = (
-            output_attentions
-            if output_attentions is not None
-            else self.config.output_attentions
-        )
+        output_attentions = (output_attentions if output_attentions is not None else self.config.vit_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
+            output_hidden_states if output_hidden_states is not None else self.config.vit_config.output_hidden_states
         )
+        return_dict = return_dict if return_dict is not None else self.config.vit_config.return_dict
 
+        # (`transpose` is done in `FlaxViTPreTrainedModel.__call__()`, so we do the same here.)
         pixel_values = jnp.transpose(pixel_values, (0, 2, 3, 1))
 
         # Handle any PRNG if needed
@@ -272,11 +252,11 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
 
     def decode(
         self,
-        input_ids,
+        decoder_input_ids,
         encoder_outputs,
         encoder_attention_mask: Optional[jnp.ndarray] = None,
-        attention_mask: Optional[jnp.ndarray] = None,
-        position_ids: Optional[jnp.ndarray] = None,
+        decoder_attention_mask: Optional[jnp.ndarray] = None,
+        decoder_position_ids: Optional[jnp.ndarray] = None,
         past_key_values: dict = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
@@ -287,29 +267,23 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
     ):
 
         output_attentions = (
-            output_attentions
-            if output_attentions is not None
-            else self.config.output_attentions
+            output_attentions if output_attentions is not None else self.config.gpt2_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
+            output_hidden_states if output_hidden_states is not None else self.config.gpt2_config.output_hidden_states
         )
+        return_dict = return_dict if return_dict is not None else self.config.gpt2_config.return_dict
 
         encoder_hidden_states = encoder_outputs[0]
         if encoder_attention_mask is None:
             batch_size, sequence_length = encoder_hidden_states.shape[:2]
             encoder_attention_mask = jnp.ones((batch_size, sequence_length))
 
-        batch_size, sequence_length = input_ids.shape
-        if attention_mask is None:
-            attention_mask = jnp.ones((batch_size, sequence_length))
+        batch_size, sequence_length = decoder_input_ids.shape
+        if decoder_attention_mask is None:
+            decoder_attention_mask = jnp.ones((batch_size, sequence_length))
 
-        if position_ids is None:
+        if decoder_position_ids is None:
             if past_key_values is not None:
                 raise ValueError(
                     "Make sure to provide `position_ids` when passing `past_key_values`."
@@ -335,26 +309,20 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
         else:
             mutable = False
 
-        def _decoder_forward(
-            module,
-            input_ids,
-            attention_mask,
-            position_ids,
-            **kwargs,
-        ):
+        def _decoder_forward(module, decoder_input_ids, decoder_attention_mask, decoder_position_ids, **kwargs):
             decoder_module = module._get_decoder_module()
             return decoder_module(
-                input_ids,
-                attention_mask,
-                position_ids,
+                decoder_input_ids,
+                decoder_attention_mask,
+                decoder_position_ids,
                 **kwargs,
             )
 
         outputs = self.module.apply(
             inputs,
-            input_ids=jnp.array(input_ids, dtype="i4"),
-            attention_mask=jnp.array(attention_mask, dtype="i4"),
-            position_ids=jnp.array(position_ids, dtype="i4"),
+            decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
+            decoder_attention_mask=jnp.array(decoder_attention_mask, dtype="i4"),
+            decoder_position_ids=jnp.array(decoder_position_ids, dtype="i4"),
             encoder_hidden_states=encoder_hidden_states,
             encoder_attention_mask=jnp.array(encoder_attention_mask, dtype="i4"),
             output_attentions=output_attentions,
@@ -380,9 +348,10 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
     def __call__(
         self,
         pixel_values: jnp.ndarray,
-        input_ids: Optional[jnp.ndarray] = None,
         attention_mask: Optional[jnp.ndarray] = None,
-        position_ids: Optional[jnp.ndarray] = None,
+        decoder_input_ids: Optional[jnp.ndarray] = None,
+        decoder_attention_mask: 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,
@@ -390,41 +359,24 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
         params: dict = None,
         dropout_rng: PRNGKey = None,
     ):
-        output_attentions = (
-            output_attentions
-            if output_attentions is not None
-            else self.config.output_attentions
-        )
+        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
+            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 (`transpose` is done in `FlaxViTPreTrainedModel.__call__()`, so we do the same here.)
         pixel_values = jnp.transpose(pixel_values, (0, 2, 3, 1))
 
-        # # prepare encoder inputs
-        # if encoder_attention_mask is None:
-        #     encoder_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:
-        #     decoder_input_ids = shift_tokens_right(
-        #         input_ids, self.config.pad_token_id, decoder_start_token_id=self.config.decoder_start_token_id
-        #     ) # TODO: Check how to use this
-
-        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(
+        if decoder_input_ids is None:
+            decoder_input_ids = self.config.decoder_start_token_id * jnp.ones((pixel_values.shape[0], 1))
+        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)
             )
 
@@ -434,9 +386,9 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
         return self.module.apply(
             {"params": params or self.params},
             pixel_values=jnp.array(pixel_values, dtype=jnp.float32),
-            input_ids=jnp.array(input_ids, dtype="i4"),
-            attention_mask=jnp.array(attention_mask, dtype="i4"),
-            position_ids=jnp.array(position_ids, dtype="i4"),
+            decoder_input_ids=jnp.array(decoder_input_ids, dtype="i4"),
+            decoder_attention_mask=jnp.array(decoder_attention_mask, 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,
@@ -445,17 +397,32 @@ class FlaxViTGPT2LMPreTrainedModel(FlaxPreTrainedModel):
         )
 
 
+# @add_start_docstrings(
+#     "The bare Bart Model transformer outputting raw hidden-states without any specific head on top.",
+#     BART_START_DOCSTRING,
+# )
+# class FlaxViTGPT2LMModel(FlaxViTGPT2LMPreTrainedModel):
+#     config: BartConfig
+#     dtype: jnp.dtype = jnp.float32  # the dtype of the computation
+#     module_class = FlaxViTGPT2LMModule
+#
+#
+# append_call_sample_docstring(
+#     FlaxBartModel, _TOKENIZER_FOR_DOC, _CHECKPOINT_FOR_DOC, FlaxSeq2SeqModelOutput, _CONFIG_FOR_DOC
+# )
+
+
 class FlaxViTGPT2LMForConditionalGeneration(FlaxViTGPT2LMPreTrainedModel):
     module_class = FlaxViTGPT2LMForConditionalGenerationModule
     dtype: jnp.dtype = jnp.float32
 
     def decode(
         self,
-        input_ids,
+        decoder_input_ids,
         encoder_outputs,
         encoder_attention_mask: Optional[jnp.ndarray] = None,
-        attention_mask: Optional[jnp.ndarray] = None,
-        position_ids: Optional[jnp.ndarray] = None,
+        decoder_attention_mask: Optional[jnp.ndarray] = None,
+        decoder_position_ids: Optional[jnp.ndarray] = None,
         past_key_values: dict = None,
         output_attentions: Optional[bool] = None,
         output_hidden_states: Optional[bool] = None,
@@ -464,135 +431,42 @@ class FlaxViTGPT2LMForConditionalGeneration(FlaxViTGPT2LMPreTrainedModel):
         params: dict = None,
         dropout_rng: PRNGKey = None,
     ):
-        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
-        )
-
-        encoder_hidden_states = encoder_outputs[0]
-        if encoder_attention_mask is None:
-            batch_size, sequence_length = encoder_hidden_states.shape[:2]
-            encoder_attention_mask = jnp.ones((batch_size, sequence_length))
 
-        batch_size, sequence_length = input_ids.shape
-        if attention_mask is None:
-            attention_mask = jnp.ones((batch_size, sequence_length))
-
-        if position_ids is None:
-            if past_key_values is not None:
-                raise ValueError(
-                    "Make sure to provide `position_ids` when passing `past_key_values`."
-                )
-
-            position_ids = jnp.broadcast_to(
-                jnp.arange(sequence_length)[None, :], (batch_size, sequence_length)
-            )
-
-        # Handle any PRNG if needed
-        rngs = {}
-        if dropout_rng is not None:
-            rngs["dropout"] = dropout_rng
-
-        inputs = {"params": params or self.params}
-
-        # if past_key_values are passed then cache is already initialized a private flag init_cache has to be
-        # passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that
-        # it can be changed by FlaxGPT2Attention module
-        if past_key_values:
-            inputs["cache"] = past_key_values
-            mutable = ["cache"]
-        else:
-            mutable = False
-
-        def _decoder_forward(
-            module,
-            input_ids,
-            attention_mask,
-            position_ids,
-            **kwargs,
-        ):
-            decoder_module = module._get_decoder_module()
-            outputs = decoder_module(
-                input_ids,
-                attention_mask,
-                position_ids,
-                **kwargs,
-            )
-            lm_logits = outputs[0]
-
-            return lm_logits, outputs
-
-        outputs = self.module.apply(
-            inputs,
-            input_ids=jnp.array(input_ids, dtype="i4"),
-            attention_mask=jnp.array(attention_mask, dtype="i4"),
-            position_ids=jnp.array(position_ids, dtype="i4"),
-            encoder_hidden_states=encoder_hidden_states,
-            encoder_attention_mask=jnp.array(encoder_attention_mask, dtype="i4"),
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-            deterministic=deterministic,
-            rngs=rngs,
-            mutable=mutable,
-            method=_decoder_forward,
+        return super().decode(
+            decoder_input_ids,
+            encoder_outputs,
+            encoder_attention_mask,
+            decoder_attention_mask,
+            decoder_position_ids,
+            past_key_values,
+            output_attentions,
+            output_hidden_states,
+            return_dict,
+            not deterministic,
+            params,
+            dropout_rng,
         )
 
-        if past_key_values is None:
-            lm_logits, outputs = outputs
-        else:
-            (lm_logits, outputs), past = outputs
-
-        if return_dict:
-            outputs = FlaxCausalLMOutputWithCrossAttentions(
-                logits=lm_logits,
-                hidden_states=outputs.decoder_hidden_states,
-                attentions=outputs.decoder_attentions,
-                cross_attentions=outputs.cross_attentions,
-            )
-        else:
-            outputs = (lm_logits,) + outputs[1:]
-
-        # add updated cache to model output
-        if past_key_values is not None and return_dict:
-            outputs["past_key_values"] = unfreeze(past["cache"])
-            return outputs
-        elif past_key_values is not None and not return_dict:
-            outputs = outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]
-
-        return outputs
-
     def prepare_inputs_for_generation(
         self,
-        input_ids,
+        decoder_input_ids,
         max_length,
-        encoder_attention_mask: Optional[jnp.DeviceArray] = None,
         attention_mask: Optional[jnp.DeviceArray] = None,
+        decoder_attention_mask: Optional[jnp.DeviceArray] = None,
         encoder_outputs=None,
         **kwargs,
     ):
         # initializing the cache
-        batch_size, seq_length = input_ids.shape
+        batch_size, seq_length = decoder_input_ids.shape
 
         past_key_values = self.init_cache(batch_size, max_length, encoder_outputs)
         # Note that usually one would have to put 0's in the attention_mask for x > input_ids.shape[-1] and x < cache_length.
         # But since the decoder uses a causal mask, those positions are masked anyways.
         # Thus we can create a single static attention_mask here, which is more efficient for compilation
         extended_attention_mask = jnp.ones((batch_size, max_length), dtype="i4")
-        if attention_mask is not None:
-            position_ids = attention_mask.cumsum(axis=-1) - 1
-            extended_attention_mask = lax.dynamic_update_slice(
-                extended_attention_mask, attention_mask, (0, 0)
-            )
+        if decoder_attention_mask is not None:
+            position_ids = decoder_attention_mask.cumsum(axis=-1) - 1
+            extended_attention_mask = lax.dynamic_update_slice(extended_attention_mask, decoder_attention_mask, (0, 0))
         else:
             position_ids = jnp.broadcast_to(
                 jnp.arange(seq_length, dtype="i4")[None, :], (batch_size, seq_length)
@@ -601,16 +475,14 @@ class FlaxViTGPT2LMForConditionalGeneration(FlaxViTGPT2LMPreTrainedModel):
         return {
             "past_key_values": past_key_values,
             "encoder_outputs": encoder_outputs,
-            "encoder_attention_mask": encoder_attention_mask,
-            "attention_mask": extended_attention_mask,
-            "position_ids": position_ids,
+            "encoder_attention_mask": attention_mask,
+            "decoder_attention_mask": extended_attention_mask,
+            "decoder_position_ids": position_ids,
         }
 
     def update_inputs_for_generation(self, model_outputs, model_kwargs):
         model_kwargs["past_key_values"] = model_outputs.past_key_values
-        model_kwargs["position_ids"] = (
-            model_kwargs["position_ids"][:, -1:] + 1
-        )
+        model_kwargs["decoder_position_ids"] = model_kwargs["decoder_position_ids"][:, -1:] + 1
         return model_kwargs
 
     @classmethod