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: Args:
      batch_size (`int`):
          batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
      max_length (`int`):
          maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
          cache.
      encoder_outputs (`Union[FlaxBaseModelOutput, tuple(tuple(jnp.ndarray)]`):
          `encoder_outputs` consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*:
          `attentions`). `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*)
          is a sequence of hidden-states at the output of the last layer of the encoder. Used in the
          cross-attention of the decoder.
docstr_tok: ["Args",":","batch_size","(","`","int","`",")",":","batch_size","used","for","fast","auto-regressive","decoding",".","Defines","the","batch","size","of","the","initialized","cache",".","max_length","(","`","int","`",")",":","maximum","possible","length","for","auto-regressive","decoding",".","Defines","the","sequence","length","of","the","initialized","cache",".","encoder_outputs","(","`","Union","[","FlaxBaseModelOutput",",","tuple","(","tuple","(","jnp.ndarray",")","]","`",")",":","`","encoder_outputs","`","consists","of","(","`","last_hidden_state","`",",","*","optional","*",":","`","hidden_states","`",",","*","optional","*",":","`","attentions","`",")",".","`","last_hidden_state","`","of","shape","`","(","batch_size",",","sequence_length",",","hidden_size",")","`",",","*","optional","*",")","is","a","sequence","of","hidden-states","at","the","output","of","the","last","layer","of","the","encoder",".","Used","in","the","cross-attention","of","the","decoder","."]
code: def init_cache(self, batch_size, max_length, encoder_outputs):
      r"""
      Args:
          batch_size (`int`):
              batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
          max_length (`int`):
              maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
              cache.
          encoder_outputs (`Union[FlaxBaseModelOutput, tuple(tuple(jnp.ndarray)]`):
              `encoder_outputs` consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*:
              `attentions`). `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*)
              is a sequence of hidden-states at the output of the last layer of the encoder. Used in the
              cross-attention of the decoder.
      """
      # 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,
          decoder_input_ids,
          decoder_attention_mask,
          decoder_position_ids,
          **kwargs,
      ):
          decoder_module = module._get_decoder_module()
          return decoder_module(
              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),
          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
      )
      return unfreeze(init_variables["cache"])
  
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 3 for calling the FlaxSpeechEncoderDecoderModel obj's underlying member method code with example usage: obj.init_cache(batch_size, max_length, encoder_outputs) and returns: unfreeze, decoder_module