class_docstr: Encoder-decoder model with T5 decoder implementing beam-search.
docstr: Predict with fast decoding beam search on a batch.
  
  This is copied and modified from T5X EncoderDecoderTransformer model in
  third_party/py/t5x/models.py.
  
  Here we refer to "parameters" for values that can be compiled into the
  model dynamically, as opposed to static configuration settings that require
  a recompile. For example, the model weights and the decoder brevity-penalty
  are parameters and can be modified without requiring a recompile. The number
  of layers, the batch size and the decoder beam size are configuration
  options that require recompilation if changed.
  
  This method can be used with a customizable decoding function as long as it
  follows the signature of `DecodeFnCallable`. In order to provide a unified
  interface for the decoding functions, we use a generic names. For example a
  beam size is a concept unique to beam search. Conceptually, it corresponds
  to the number of sequences returned by the beam search.  Therefore, the
  generic argument `num_decodes` corresponds to the beam size if
  `decode_fn` is a beam search. For temperature sampling, `num_decodes`
  corresponds to the number of indepedent sequences to be sampled. Typically
  `num_decodes = 1` is used for tempeature sampling.
  
  If `return_all_decodes = True`, the return tuple contains the predictions
  with a shape [batch, num_decodes, max_decode_len] and the scores (i.e., log
  probability of the generated sequence) with a shape [batch, num_decodes].
  
  If `return_all_decodes = False`, the return tuple contains the predictions
  with a shape [batch, max_decode_len] and the scores with a shape [batch].
  
  `decoder_params` can be used to pass dynamic configurations to
  `decode_fn`. An example usage is to pass different random seed (i.e.,
  `jax.random.PRNGKey(seed)` with different `seed` value). This can be done by
  setting `decoder_params['decode_rng'] = jax.random.PRNGKey(seed)`.
  
  Args:
    params: model parameters.
    batch: a batch of inputs.
    decode_fn: function implementing the decode method.
    eos_id: EOS token id in the vocabulary.
    decoder_params: additional (model-independent) parameters for the decoder.
    return_all_decodes: whether to return the entire beam or just the top-1.
    num_decodes: the number of beams to use in beam search.
    alpha: length penalty factor for beam search.
    decoding_method: decoding method.
    temperature: temperature for nucleus sampling.
    vocabulary_size: size of the vocabulary for the textual tokens.
  
  Returns:
    A tuple containing:
      the batch of predictions, with the entire beam if requested
      an auxiliary dictionary of decoder scores
docstr_tok: ["Predict","with","fast","decoding","beam","search","on","a","batch",".","This","is","copied","and","modified","from","T5X","EncoderDecoderTransformer","model","in","third_party\/py\/t5x\/models.py",".","Here","we","refer","to","``","parameters","''","for","values","that","can","be","compiled","into","the","model","dynamically",",","as","opposed","to","static","configuration","settings","that","require","a","recompile",".","For","example",",","the","model","weights","and","the","decoder","brevity-penalty","are","parameters","and","can","be","modified","without","requiring","a","recompile",".","The","number","of","layers",",","the","batch","size","and","the","decoder","beam","size","are","configuration","options","that","require","recompilation","if","changed",".","This","method","can","be","used","with","a","customizable","decoding","function","as","long","as","it","follows","the","signature","of","`","DecodeFnCallable","`",".","In","order","to","provide","a","unified","interface","for","the","decoding","functions",",","we","use","a","generic","names",".","For","example","a","beam","size","is","a","concept","unique","to","beam","search",".","Conceptually",",","it","corresponds","to","the","number","of","sequences","returned","by","the","beam","search",".","Therefore",",","the","generic","argument","`","num_decodes","`","corresponds","to","the","beam","size","if","`","decode_fn","`","is","a","beam","search",".","For","temperature","sampling",",","`","num_decodes","`","corresponds","to","the","number","of","indepedent","sequences","to","be","sampled",".","Typically","`","num_decodes","=","1","`","is","used","for","tempeature","sampling",".","If","`","return_all_decodes","=","True","`",",","the","return","tuple","contains","the","predictions","with","a","shape","[","batch",",","num_decodes",",","max_decode_len","]","and","the","scores","(","i.e.",",","log","probability","of","the","generated","sequence",")","with","a","shape","[","batch",",","num_decodes","]",".","If","`","return_all_decodes","=","False","`",",","the","return","tuple","contains","the","predictions","with","a","shape","[","batch",",","max_decode_len","]","and","the","scores","with","a","shape","[","batch","]",".","`","decoder_params","`","can","be","used","to","pass","dynamic","configurations","to","`","decode_fn","`",".","An","example","usage","is","to","pass","different","random","seed","(","i.e.",",","`","jax.random.PRNGKey","(","seed",")","`","with","different","`","seed","`","value",")",".","This","can","be","done","by","setting","`","decoder_params","[","'decode_rng","'","]","=","jax.random.PRNGKey","(","seed",")","`",".","Args",":","params",":","model","parameters",".","batch",":","a","batch","of","inputs",".","decode_fn",":","function","implementing","the","decode","method",".","eos_id",":","EOS","token","id","in","the","vocabulary",".","decoder_params",":","additional","(","model-independent",")","parameters","for","the","decoder",".","return_all_decodes",":","whether","to","return","the","entire","beam","or","just","the","top-1",".","num_decodes",":","the","number","of","beams","to","use","in","beam","search",".","alpha",":","length","penalty","factor","for","beam","search",".","decoding_method",":","decoding","method",".","temperature",":","temperature","for","nucleus","sampling",".","vocabulary_size",":","size","of","the","vocabulary","for","the","textual","tokens",".","Returns",":","A","tuple","containing",":","the","batch","of","predictions",",","with","the","entire","beam","if","requested","an","auxiliary","dictionary","of","decoder","scores"]
code: def predict_batch_with_aux(
      self,
      params: PyTree,
      batch: PyTree,
      decode_fn: DecodeFnCallable,
      eos_id: int = 1,
      decoder_params: Optional[Dict[str, Any]] = None,
      return_all_decodes: bool = False,
      num_decodes: int = 1,
      alpha: float = 0.6,
      decoding_method: str = "beamsearch",
      temperature: float = 1.0,
      vocabulary_size: int = 32128,
  ):
      """Predict with fast decoding beam search on a batch.
  
      This is copied and modified from T5X EncoderDecoderTransformer model in
      third_party/py/t5x/models.py.
  
      Here we refer to "parameters" for values that can be compiled into the
      model dynamically, as opposed to static configuration settings that require
      a recompile. For example, the model weights and the decoder brevity-penalty
      are parameters and can be modified without requiring a recompile. The number
      of layers, the batch size and the decoder beam size are configuration
      options that require recompilation if changed.
  
      This method can be used with a customizable decoding function as long as it
      follows the signature of `DecodeFnCallable`. In order to provide a unified
      interface for the decoding functions, we use a generic names. For example a
      beam size is a concept unique to beam search. Conceptually, it corresponds
      to the number of sequences returned by the beam search.  Therefore, the
      generic argument `num_decodes` corresponds to the beam size if
      `decode_fn` is a beam search. For temperature sampling, `num_decodes`
      corresponds to the number of indepedent sequences to be sampled. Typically
      `num_decodes = 1` is used for tempeature sampling.
  
      If `return_all_decodes = True`, the return tuple contains the predictions
      with a shape [batch, num_decodes, max_decode_len] and the scores (i.e., log
      probability of the generated sequence) with a shape [batch, num_decodes].
  
      If `return_all_decodes = False`, the return tuple contains the predictions
      with a shape [batch, max_decode_len] and the scores with a shape [batch].
  
      `decoder_params` can be used to pass dynamic configurations to
      `decode_fn`. An example usage is to pass different random seed (i.e.,
      `jax.random.PRNGKey(seed)` with different `seed` value). This can be done by
      setting `decoder_params['decode_rng'] = jax.random.PRNGKey(seed)`.
  
      Args:
        params: model parameters.
        batch: a batch of inputs.
        decode_fn: function implementing the decode method.
        eos_id: EOS token id in the vocabulary.
        decoder_params: additional (model-independent) parameters for the decoder.
        return_all_decodes: whether to return the entire beam or just the top-1.
        num_decodes: the number of beams to use in beam search.
        alpha: length penalty factor for beam search.
        decoding_method: decoding method.
        temperature: temperature for nucleus sampling.
        vocabulary_size: size of the vocabulary for the textual tokens.
  
      Returns:
        A tuple containing:
          the batch of predictions, with the entire beam if requested
          an auxiliary dictionary of decoder scores
      """
      # Prepare zeroed-out autoregressive cache.
      encoder_inputs = jax.tree_map(
          jnp.ones_like, batch["encoder_inputs"]
      )
      decoder_inputs = jax.tree_map(
          jnp.ones_like, batch["decoder_inputs"]
      )
      _, variables_with_cache = self.flax_model.apply(
          params,
          encoder_inputs,
          decoder_inputs,
          decode=True,
          train=False,
          mutable=["cache"],
      )
      cache = variables_with_cache["cache"]
  
      # 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
      )
      non_expanded_encoded = self.flax_model.apply(
          params,
          batch["encoder_inputs"],
          train=False,
          method=self.flax_model.encode,
      )
      encoded_inputs = jax.tree_map(
          beam_expand_fn, non_expanded_encoded
      )
      if isinstance(encoded_inputs, dict):  # set decoder mask
          batch["decoder_inputs"][
              "encoder_input_tokens"
          ] = encoded_inputs["mask"]
          encoded_inputs = encoded_inputs["encoded"]
  
      # Set the all output embeddings to be valid inputs if encoder_input_tokens
      # are not provided. Note that this tensor should be beam-extended too.
      decoder_inputs = batch["decoder_inputs"]
      if "encoder_input_tokens" not in decoder_inputs:
          input_masks = jnp.ones(encoded_inputs.shape[:-1])
      else:
          input_masks = decoder_inputs["encoder_input_tokens"]
  
      tokens_ids_to_logits = functools.partial(
          self._compute_logits_from_slice,
          all_variables=params,
          encoded_inputs=encoded_inputs,
          input_masks=input_masks,
          max_decode_length=decoder_inputs[
              "decoder_input_tokens"
          ].shape[1],
      )
  
      if decoder_params is None:
          decoder_params = {}
  
      # `decoder_prompt_inputs` is only used to obtain batch size
      # and max decode length information here.
      decoder_prompt_inputs = jnp.zeros_like(
          decoder_inputs["decoder_input_tokens"]
      )
  
      # Using the above-defined single-step decoder function, run a
      # beam search over possible sequences given input encoding.
      # decodes: [batch, num_decodes, max_decode_len + 1]
      # scores: [batch, num_decodes]
      if decoding_method == "temperature_sample":
          decodes, scores = decode_fn(
              inputs=decoder_prompt_inputs,
              cache=cache,
              tokens_to_logits=tokens_ids_to_logits,
              eos_id=eos_id,
              topp=alpha,
              topk=0,
              temperature=temperature,
              num_decodes=num_decodes,
              cache_offset=0,
              **decoder_params,
          )
      else:  # beam search
          decodes, scores = decode_fn(
              inputs=decoder_prompt_inputs,
              cache=cache,
              tokens_to_logits=tokens_ids_to_logits,
              eos_id=eos_id,
              alpha=alpha,
              num_decodes=num_decodes,
              cache_offset=0,
              **decoder_params,
          )
  
      # Beam search returns [n_batch, n_beam, n_length] with beam dimension sorted
      # in increasing order of log-probability.
      # Return the highest scoring beam sequence.
      if return_all_decodes:
          return decodes, {"scores": scores}
      else:
          return decodes[:, -1, :], {"scores": scores[:, -1]}
  
imports: import functools
  from typing import Any, Dict, Mapping, Optional, Tuple
  from absl import logging
  import flax.linen
  from flax.training import common_utils
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.baselines import vit
  from scenic.projects.t5 import layers
  from scenic.projects.t5 import model
  from t5x import decoding
  from t5x.models import DecodeFnCallable
label_desc: Use image node_id 2 for calling the EncoderWithT5DecoderModel obj's underlying member method code with example usage: obj.predict_batch_with_aux(params, batch, decode_fn, eos_id, decoder_params, return_all_decodes, num_decodes, alpha, decoding_method, temperature, vocabulary_size) and returns: decodes,  dict, decodes,  dict