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: Single autoregressive decode step with caching.
docstr_tok: ["Single","autoregressive","decode","step","with","caching","."]
code: def single_decode_step(
      self,
      decoding_state: decoding.DecodingState,
      variables: constants.PyTree,
      encoded_inputs: jnp.ndarray,
      input_masks: jnp.ndarray,
      max_decode_length: int,
  ) -> Tuple[jnp.ndarray, Mapping[str, jnp.ndarray]]:
      """Single autoregressive decode step with caching."""
      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, **variables},
          encoded=encoded_inputs,
          encoder_input_tokens=input_masks,
          decoder_input_tokens=flat_ids,
          decoder_target_tokens=flat_ids,
          decode=True,
          enable_dropout=False,
          max_decode_length=max_decode_length,
          mutable=["cache"],
          method=self.flax_model.t5_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, 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 6 for calling the KnowledgeFIDModel obj's underlying member method code with example usage: obj.single_decode_step(decoding_state, variables, encoded_inputs, input_masks, max_decode_length) and returns: flat_logits,  new_flat_cache