class_docstr: FID model with a retrieval module over a knowledge memory.
class_docstr_tok: ["FID","model","with","a","retrieval","module","over","a","knowledge","memory","."]
docstr: Apply inference with autoregressive decoding.
  
  Apply t5x autoregressive decoding with cache using either their
  beam_search or temperature_sample decoding technique.
  
  Args:
    variables: variables of the models.
    decoder_input_tokens:  # B×O.
    decoder_target_tokens:  # B×O.
    encoder_input_image:  # B×W×H×3.
    encoder_input_tokens: # B×I.
    num_decodes: number of outputs generated per input for the decode search.
    debug: Whether in debug mode or not.
    beam_search: If True, do beam search. If False, do temperature sampling.
    decoder_params: Additional decoding parameters. These provide additional
      parameters to beam_search or temperature_sample (see decoder module).
    return_all_decodes: If True, return all decodes. Otherwise only return the
      top scored decoding.
    use_memory: whether use on-device memory.
    retrieve_local: whether only retrieve in local host or across hosts.
    **args: other possible arguments.
  
  Returns:
    logits array from the final decoder.
docstr_tok: ["Apply","inference","with","autoregressive","decoding",".","Apply","t5x","autoregressive","decoding","with","cache","using","either","their","beam_search","or","temperature_sample","decoding","technique",".","Args",":","variables",":","variables","of","the","models",".","decoder_input_tokens",":","#","B\u00d7O",".","decoder_target_tokens",":","#","B\u00d7O",".","encoder_input_image",":","#","B\u00d7W\u00d7H\u00d73",".","encoder_input_tokens",":","#","B\u00d7I",".","num_decodes",":","number","of","outputs","generated","per","input","for","the","decode","search",".","debug",":","Whether","in","debug","mode","or","not",".","beam_search",":","If","True",",","do","beam","search",".","If","False",",","do","temperature","sampling",".","decoder_params",":","Additional","decoding","parameters",".","These","provide","additional","parameters","to","beam_search","or","temperature_sample","(","see","decoder","module",")",".","return_all_decodes",":","If","True",",","return","all","decodes",".","Otherwise","only","return","the","top","scored","decoding",".","use_memory",":","whether","use","on-device","memory",".","retrieve_local",":","whether","only","retrieve","in","local","host","or","across","hosts",".","*","*","args",":","other","possible","arguments",".","Returns",":","logits","array","from","the","final","decoder","."]
code: def apply_with_autoregressive_decoding(
      self,
      variables: constants.PyTree,
      decoder_input_tokens: jnp.ndarray,
      decoder_target_tokens: jnp.ndarray,
      encoder_input_image: Optional[jnp.ndarray] = None,
      encoder_input_tokens: Optional[jnp.ndarray] = None,
      num_decodes: int = 1,
      debug: bool = False,
      beam_search: bool = True,
      decoder_params: Optional[dict[str, Any]] = None,
      return_all_decodes: bool = False,
      use_memory=False,
      retrieve_local=False,
      **args,
  ):
      """Apply inference with autoregressive decoding.
  
      Apply t5x autoregressive decoding with cache using either their
      beam_search or temperature_sample decoding technique.
  
      Args:
        variables: variables of the models.
        decoder_input_tokens:  # B×O.
        decoder_target_tokens:  # B×O.
        encoder_input_image:  # B×W×H×3.
        encoder_input_tokens: # B×I.
        num_decodes: number of outputs generated per input for the decode search.
        debug: Whether in debug mode or not.
        beam_search: If True, do beam search. If False, do temperature sampling.
        decoder_params: Additional decoding parameters. These provide additional
          parameters to beam_search or temperature_sample (see decoder module).
        return_all_decodes: If True, return all decodes. Otherwise only return the
          top scored decoding.
        use_memory: whether use on-device memory.
        retrieve_local: whether only retrieve in local host or across hosts.
        **args: other possible arguments.
  
      Returns:
        logits array from the final decoder.
      """
      # Prepare zeroed-out autoregressive cache.
      _, model_state_with_cache = self.flax_model.apply(
          variables=variables,
          encoder_input_image=encoder_input_image,
          encoder_input_tokens=encoder_input_tokens,
          decoder_input_tokens=decoder_input_tokens,
          decoder_target_tokens=decoder_target_tokens,
          train=False,
          only_encode=False,
          decode=True,
          mutable=["cache"],
          debug=debug,
          use_memory=use_memory,
          retrieve_local=retrieve_local,
      )
  
      # Call model to get the features consumed by the decoder. Skip the
      # the decoding part itself.
      out_dict = self.flax_model.apply(
          variables=variables,
          encoder_input_image=encoder_input_image,
          encoder_input_tokens=encoder_input_tokens,
          decoder_input_tokens=decoder_input_tokens,
          decoder_target_tokens=decoder_target_tokens,
          train=False,
          only_encode=True,
          debug=debug,
          use_memory=use_memory,
          retrieve_local=retrieve_local,
      )
      retr_top_image = out_dict["retr_data"]["image"][:, 0]
  
      # Prepare transformer fast-decoder call for beam search: for beam search, we
      # need to set up our decoder model to handle a batch size equal to
      # batch_size * num_decodes, where each batch item's data is expanded
      # in-place rather than tiled.
      # i.e. if we denote each batch element subtensor as el[n]:
      # [el0, el1, el2] --> beamsize=2 --> [el0,el0,el1,el1,el2,el2]
      # [batch * num_decodes, input_len, emb_dim]
      beam_expand_fn = functools.partial(
          decoding.flat_batch_beam_expand, beam_size=num_decodes
      )
      encoded_inputs = jax.tree_map(
          beam_expand_fn, out_dict["fused_emb"]
      )
      encoded_masks = jax.tree_util.tree_map(
          beam_expand_fn, out_dict["fused_mask"]
      )
      bsz = decoder_input_tokens.shape[0]
      max_decode_length = decoder_input_tokens.shape[-1]
      # Define the token2logit function for a single decoding step.
      tokens_ids_to_logits = functools.partial(
          self.single_decode_step,
          variables=variables,
          encoded_inputs=encoded_inputs,
          input_masks=encoded_masks,
          max_decode_length=decoder_input_tokens.shape[-1],
      )
  
      if decoder_params is None:
          decoder_params = {}
      # For beam search, `decoder_prompt_inputs` is only used to obtain batch size
      # and max decode length information. For temperature sampling,
      # `decod_prompt_inputs` will be filled with the sampled ids.
      decoder_prompt_inputs = jnp.zeros([bsz, max_decode_length - 1])
      bos_inputs = jnp.ones([bsz, 1]) * data_utils.BOS_ID
      decoder_prompt_inputs = jnp.concatenate(
          (bos_inputs, decoder_prompt_inputs), axis=-1, dtype=jnp.int32
      )
      if beam_search:
          decodes, scores = decoding.beam_search(
              inputs=decoder_prompt_inputs,
              cache=model_state_with_cache["cache"],
              tokens_to_logits=tokens_ids_to_logits,
              eos_id=data_utils.EOS_ID,
              num_decodes=num_decodes,
              cache_offset=0,
              **decoder_params,
          )
      else:
          decodes, scores = decoding.temperature_sample(
              inputs=decoder_prompt_inputs,
              cache=model_state_with_cache["cache"],
              tokens_to_logits=tokens_ids_to_logits,
              eos_id=data_utils.EOS_ID,
              num_decodes=num_decodes,
              cache_offset=0,
              initial_index=jnp.zeros([bsz], dtype=jnp.int32),
              **decoder_params,
          )
      if return_all_decodes:
          return decodes, scores, retr_top_image
      else:
          return decodes[:, -1, :], scores[:, -1], retr_top_image
  
imports: import functools
  from typing import Any, Dict, Mapping, Optional, Tuple, List
  from absl import logging
  import flax.linen
  import jax
  from jax.experimental import host_callback
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.projects.knowledge_visual_language.data import data_utils
  from scenic.projects.knowledge_visual_language.models import constants
  from scenic.projects.knowledge_visual_language.models import fusion_in_decoder_soft
  from scenic.projects.knowledge_visual_language.models import layers
  from scenic.projects.knowledge_visual_language.models import local_memory
  from scenic.projects.knowledge_visual_language.models import losses
  from scenic.projects.knowledge_visual_language.models import metrics
  from t5x import decoding
label_desc: Use image node_id 7 for calling the KnowledgeFIDModel obj's underlying member method code with example usage: obj.apply_with_autoregressive_decoding(variables, decoder_input_tokens, decoder_target_tokens, encoder_input_image, encoder_input_tokens, num_decodes, debug, beam_search, decoder_params, return_all_decodes, use_memory, retrieve_local) and returns: decodes,  scores,  retr_top_image, decodes,  scores,  retr_top_image