class_docstr: Encoder-decoder model with T5 decoder implementing beam-search.
docstr: Token slice to logits from decoder model.
docstr_tok: ["Token","slice","to","logits","from","decoder","model","."]
code: def _compute_logits_from_slice(
      self,
      decoding_state: decoding.DecodingState,
      all_variables: PyTree,
      encoded_inputs: jnp.ndarray,
      input_masks: jnp.ndarray,
      max_decode_length: int,
  ) -> Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]:
      """Token slice to logits from decoder model."""
      flat_ids = decoding_state.cur_token
      flat_cache = decoding_state.cache
      # flat_ids: [batch * beam, seq_len=1]
      # cache is expanded inside beam_search to become flat_cache
      # flat_cache: [batch * beam, num_heads, depth_per_head, max_decode_len]
      # flat_logits: [batch * beam, seq_len=1, vocab]
      flat_logits, new_vars = self.flax_model.apply(
          {"cache": flat_cache, **all_variables},
          encoded_inputs,
          {
              "encoder_input_tokens": input_masks,
              "decoder_input_tokens": flat_ids,
              "decoder_target_tokens": flat_ids,
          },
          train=False,
          decode=True,
          max_decode_length=max_decode_length,
          mutable=["cache"],
          method=self.flax_model.decode,
      )
      # Remove sequence length dimension since it's always 1 during decoding.
      flat_logits = jnp.squeeze(flat_logits, axis=1)
      new_flat_cache = new_vars["cache"]
      return flat_logits, new_flat_cache
  
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 1 for calling the EncoderWithT5DecoderModel obj's underlying member method code with example usage: obj._compute_logits_from_slice(decoding_state, all_variables, encoded_inputs, input_masks, max_decode_length) and returns: flat_logits,  new_flat_cache