class_docstr: [`FlaxVisionEncoderDecoderModel`] is a generic model class that will be instantiated as a transformer architecture
  with the module (flax.nn.Module) of one of the base vision model classes of the library as encoder module and
  another one as decoder module when created with the :meth*~transformers.FlaxAutoModel.from_pretrained* class method
  for the encoder and :meth*~transformers.FlaxAutoModelForCausalLM.from_pretrained* class method for the decoder.
class_docstr_tok: ["[","`","FlaxVisionEncoderDecoderModel","`","]","is","a","generic","model","class","that","will","be","instantiated","as","a","transformer","architecture","with","the","module","(","flax.nn.Module",")","of","one","of","the","base","vision","model","classes","of","the","library","as","encoder","module","and","another","one","as","decoder","module","when","created","with","the",":","meth","*","~transformers.FlaxAutoModel.from_pretrained","*","class","method","for","the","encoder","and",":","meth","*","~transformers.FlaxAutoModelForCausalLM.from_pretrained","*","class","method","for","the","decoder","."]
docstr: Returns:
  
  Example:
  
  ```python
  >>> from transformers import ViTFeatureExtractor, FlaxVisionEncoderDecoderModel
  >>> import jax.numpy as jnp
  >>> from PIL import Image
  >>> import requests
  
  >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
  >>> image = Image.open(requests.get(url, stream=True).raw)
  
  >>> feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
  
  >>> # initialize a vit-gpt2 from pretrained ViT and GPT2 models. Note that the cross-attention layers will be randomly initialized
  >>> model = FlaxVisionEncoderDecoderModel.from_encoder_decoder_pretrained(
  ...     "google/vit-base-patch16-224-in21k", "gpt2"
  ... )
  
  >>> pixel_values = feature_extractor(images=image, return_tensors="np").pixel_values
  >>> encoder_outputs = model.encode(pixel_values)
  
  >>> decoder_start_token_id = model.config.decoder.bos_token_id
  >>> decoder_input_ids = jnp.ones((pixel_values.shape[0], 1), dtype="i4") * decoder_start_token_id
  
  >>> outputs = model.decode(decoder_input_ids, encoder_outputs)
  >>> logits = outputs.logits
  ```
docstr_tok: ["Returns",":","Example",":","``","`","python",">",">",">","from","transformers","import","ViTFeatureExtractor",",","FlaxVisionEncoderDecoderModel",">",">",">","import","jax.numpy","as","jnp",">",">",">","from","PIL","import","Image",">",">",">","import","requests",">",">",">","url","=","``","http",":","\/\/images.cocodataset.org\/val2017\/000000039769.jpg","''",">",">",">","image","=","Image.open","(","requests.get","(","url",",","stream=True",")",".raw",")",">",">",">","feature_extractor","=","ViTFeatureExtractor.from_pretrained","(","``","google\/vit-base-patch16-224-in21k","''",")",">",">",">","#","initialize","a","vit-gpt2","from","pretrained","ViT","and","GPT2","models",".","Note","that","the","cross-attention","layers","will","be","randomly","initialized",">",">",">","model","=","FlaxVisionEncoderDecoderModel.from_encoder_decoder_pretrained","(","...","``","google\/vit-base-patch16-224-in21k","''",",","``","gpt2","''","...",")",">",">",">","pixel_values","=","feature_extractor","(","images=image",",","return_tensors=","''","np","''",")",".pixel_values",">",">",">","encoder_outputs","=","model.encode","(","pixel_values",")",">",">",">","decoder_start_token_id","=","model.config.decoder.bos_token_id",">",">",">","decoder_input_ids","=","jnp.ones","(","(","pixel_values.shape","[","0","]",",","1",")",",","dtype=","''","i4","''",")","*","decoder_start_token_id",">",">",">","outputs","=","model.decode","(","decoder_input_ids",",","encoder_outputs",")",">",">",">","logits","=","outputs.logits","``","`"]
code: def decode(
      self,
      decoder_input_ids,
      encoder_outputs,
      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,
      return_dict: Optional[bool] = None,
      train: bool = False,
      params: dict = None,
      dropout_rng: PRNGKey = None,
  ):
      r"""
      Returns:
  
      Example:
  
      ```python
      >>> from transformers import ViTFeatureExtractor, FlaxVisionEncoderDecoderModel
      >>> import jax.numpy as jnp
      >>> from PIL import Image
      >>> import requests
  
      >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
      >>> image = Image.open(requests.get(url, stream=True).raw)
  
      >>> feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")
  
      >>> # initialize a vit-gpt2 from pretrained ViT and GPT2 models. Note that the cross-attention layers will be randomly initialized
      >>> model = FlaxVisionEncoderDecoderModel.from_encoder_decoder_pretrained(
      ...     "google/vit-base-patch16-224-in21k", "gpt2"
      ... )
  
      >>> pixel_values = feature_extractor(images=image, return_tensors="np").pixel_values
      >>> encoder_outputs = model.encode(pixel_values)
  
      >>> decoder_start_token_id = model.config.decoder.bos_token_id
      >>> decoder_input_ids = jnp.ones((pixel_values.shape[0], 1), dtype="i4") * decoder_start_token_id
  
      >>> outputs = model.decode(decoder_input_ids, encoder_outputs)
      >>> logits = outputs.logits
      ```"""
      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]
  
      batch_size, sequence_length = encoder_hidden_states.shape[:2]
      encoder_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 decoder_position_ids is None:
          if past_key_values is not None:
              raise ValueError(
                  "Make sure to provide `decoder_position_ids` when passing `past_key_values`."
              )
  
          decoder_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 FlaxBartAttention module
      if past_key_values:
          inputs["cache"] = past_key_values
          mutable = ["cache"]
      else:
          mutable = False
  
      def _decoder_forward(
          module,
          decoder_input_ids,
          decoder_attention_mask,
          decoder_position_ids,
          encoder_hidden_states,
          **kwargs,
      ):
          projection_module = module._get_projection_module()
          decoder_module = module._get_decoder_module()
  
          # optionally project encoder_hidden_states
          if projection_module is not None:
              encoder_hidden_states = projection_module(
                  encoder_hidden_states
              )
  
          return decoder_module(
              decoder_input_ids,
              decoder_attention_mask,
              decoder_position_ids,
              encoder_hidden_states,
              **kwargs,
          )
  
      outputs = self.module.apply(
          inputs,
          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,
          output_hidden_states=output_hidden_states,
          return_dict=return_dict,
          deterministic=not train,
          rngs=rngs,
          mutable=mutable,
          method=_decoder_forward,
      )
  
      # add updated cache to model output
      if past_key_values is not None and return_dict:
          outputs, past = outputs
          outputs["past_key_values"] = unfreeze(past["cache"])
          return outputs
      elif past_key_values is not None and not return_dict:
          outputs, past = outputs
          outputs = (
              outputs[:1] + (unfreeze(past["cache"]),) + outputs[1:]
          )
  
      return outputs
  
imports: import os
  from typing import Optional, Tuple, Union
  import flax.linen
  import jax
  import jax.numpy
  from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
  from flax.traverse_util import flatten_dict, unflatten_dict
  from jax import lax
  from jax.random import PRNGKey
  from ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxCausalLMOutputWithCrossAttentions, FlaxSeq2SeqLMOutput
  from ...modeling_flax_utils import FlaxPreTrainedModel
  from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
  from ..auto.configuration_auto import AutoConfig
  from ..auto.modeling_flax_auto import FlaxAutoModel, FlaxAutoModelForCausalLM
  from .configuration_vision_encoder_decoder import VisionEncoderDecoderConfig
label_desc: Use image node_id 5 for calling the FlaxVisionEncoderDecoderModel obj's underlying member method code with example usage: obj.decode(decoder_input_ids, encoder_outputs, decoder_attention_mask, decoder_position_ids, past_key_values, output_attentions, output_hidden_states, return_dict, train, params, dropout_rng) and returns: outputs, decoder_module, outputs