class_docstr: [`FlaxSpeechEncoderDecoderModel`] is a generic model class that will be instantiated as a transformer architecture
  with the module (flax.nn.Module) of one of the base 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: ["[","`","FlaxSpeechEncoderDecoderModel","`","]","is","a","generic","model","class","that","will","be","instantiated","as","a","transformer","architecture","with","the","module","(","flax.nn.Module",")","of","one","of","the","base","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 FlaxSpeechEncoderDecoderModel
  >>> import jax.numpy as jnp
  
  >>> # initialize a wav2vec2-2-bart from pretrained wav2vec2 and bart models. Note that the cross-attention layers will be randomly initialized
  >>> model = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained(
  ...     "facebook/wav2vec2-large-lv60", "facebook/bart-large"
  ... )
  
  >>> inputs = jnp.ones((2, 5000), dtype=jnp.float32)
  >>> encoder_outputs = model.encode(inputs)
  
  >>> decoder_start_token_id = model.config.decoder.bos_token_id
  >>> decoder_input_ids = jnp.ones((inputs.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","FlaxSpeechEncoderDecoderModel",">",">",">","import","jax.numpy","as","jnp",">",">",">","#","initialize","a","wav2vec2-2-bart","from","pretrained","wav2vec2","and","bart","models",".","Note","that","the","cross-attention","layers","will","be","randomly","initialized",">",">",">","model","=","FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained","(","...","``","facebook\/wav2vec2-large-lv60","''",",","``","facebook\/bart-large","''","...",")",">",">",">","inputs","=","jnp.ones","(","(","2",",","5000",")",",","dtype=jnp.float32",")",">",">",">","encoder_outputs","=","model.encode","(","inputs",")",">",">",">","decoder_start_token_id","=","model.config.decoder.bos_token_id",">",">",">","decoder_input_ids","=","jnp.ones","(","(","inputs.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,
      encoder_attention_mask: 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,
      return_dict: Optional[bool] = None,
      train: bool = False,
      params: dict = None,
      dropout_rng: PRNGKey = None,
  ):
      r"""
      Returns:
  
      Example:
  
      ```python
      >>> from transformers import FlaxSpeechEncoderDecoderModel
      >>> import jax.numpy as jnp
  
      >>> # initialize a wav2vec2-2-bart from pretrained wav2vec2 and bart models. Note that the cross-attention layers will be randomly initialized
      >>> model = FlaxSpeechEncoderDecoderModel.from_encoder_decoder_pretrained(
      ...     "facebook/wav2vec2-large-lv60", "facebook/bart-large"
      ... )
  
      >>> inputs = jnp.ones((2, 5000), dtype=jnp.float32)
      >>> encoder_outputs = model.encode(inputs)
  
      >>> decoder_start_token_id = model.config.decoder.bos_token_id
      >>> decoder_input_ids = jnp.ones((inputs.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]
      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 = 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
  
      params = {"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:
          params["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=encoder_hidden_states,
              **kwargs,
          )
  
      outputs = self.module.apply(
          params,
          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_speech_encoder_decoder import SpeechEncoderDecoderConfig
label_desc: Use image node_id 6 for calling the FlaxSpeechEncoderDecoderModel obj's underlying member method code with example usage: obj.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, train, params, dropout_rng) and returns: outputs, decoder_module, outputs