class_docstr: FID model (https://arxiv.org/pdf/2007.01282.pdf) with a retrieval module over a knowledge memory.
class_docstr_tok: ["FID","model","(","https",":","\/\/arxiv.org\/pdf\/2007.01282.pdf",")","with","a","retrieval","module","over","a","knowledge","memory","."]
docstr: Conduct online retrieval and retrieval-augmented generataion.
  
  Args:
    decoder_input_tokens:  # B×O.
    decoder_target_tokens:  # B×O.
    encoder_input_image:  # B×W×H×3.
    encoder_input_tokens: # B×I.
    retr_texts:  # B×K×L.
    retr_images:  # B×K×W×H×3.
    device_id: index of TPU device.
    train: whether using train mode.
    decode: whether in decode mode.
    max_decode_length: maximum decode token length.
    use_memory: whether use on-device memory.
    use_psudo_retr: whether to use psudo retrieved groundtruth for guidance.
    retrieve_local: whether only retrieve in local host or across hosts.
    no_memory: whether not using any retrieval.
    debug: whether use debug mode.
    frozen_base: whether froze the whole encoder.
    only_encode: skip decoding and only return encoded tokens.
    **args: other possible arguments.
  
  Returns:
    output dictionary containing final and intermediate results.
docstr_tok: ["Conduct","online","retrieval","and","retrieval-augmented","generataion",".","Args",":","decoder_input_tokens",":","#","B\u00d7O",".","decoder_target_tokens",":","#","B\u00d7O",".","encoder_input_image",":","#","B\u00d7W\u00d7H\u00d73",".","encoder_input_tokens",":","#","B\u00d7I",".","retr_texts",":","#","B\u00d7K\u00d7L",".","retr_images",":","#","B\u00d7K\u00d7W\u00d7H\u00d73",".","device_id",":","index","of","TPU","device",".","train",":","whether","using","train","mode",".","decode",":","whether","in","decode","mode",".","max_decode_length",":","maximum","decode","token","length",".","use_memory",":","whether","use","on-device","memory",".","use_psudo_retr",":","whether","to","use","psudo","retrieved","groundtruth","for","guidance",".","retrieve_local",":","whether","only","retrieve","in","local","host","or","across","hosts",".","no_memory",":","whether","not","using","any","retrieval",".","debug",":","whether","use","debug","mode",".","frozen_base",":","whether","froze","the","whole","encoder",".","only_encode",":","skip","decoding","and","only","return","encoded","tokens",".","*","*","args",":","other","possible","arguments",".","Returns",":","output","dictionary","containing","final","and","intermediate","results","."]
code: def __call__(
      self,
      decoder_input_tokens,
      decoder_target_tokens,
      encoder_input_image=None,
      encoder_input_tokens=None,
      retr_texts=None,
      retr_images=None,
      device_id=0,
      train=False,
      decode=False,
      max_decode_length=None,
      use_memory=False,
      use_psudo_retr=False,
      retrieve_local=False,
      no_memory=False,
      debug=False,
      frozen_base=True,
      only_encode=False,
      **args,
  ):
      """Conduct online retrieval and retrieval-augmented generataion.
  
      Args:
        decoder_input_tokens:  # B×O.
        decoder_target_tokens:  # B×O.
        encoder_input_image:  # B×W×H×3.
        encoder_input_tokens: # B×I.
        retr_texts:  # B×K×L.
        retr_images:  # B×K×W×H×3.
        device_id: index of TPU device.
        train: whether using train mode.
        decode: whether in decode mode.
        max_decode_length: maximum decode token length.
        use_memory: whether use on-device memory.
        use_psudo_retr: whether to use psudo retrieved groundtruth for guidance.
        retrieve_local: whether only retrieve in local host or across hosts.
        no_memory: whether not using any retrieval.
        debug: whether use debug mode.
        frozen_base: whether froze the whole encoder.
        only_encode: skip decoding and only return encoded tokens.
        **args: other possible arguments.
  
      Returns:
        output dictionary containing final and intermediate results.
      """
      bsz = decoder_input_tokens.shape[0]
  
      out_dict = {}
      base_vals, base_masks, base_query = self.encode_query(
          encoder_input_image=encoder_input_image,
          encoder_input_tokens=encoder_input_tokens,
          frozen_base=frozen_base,
      )
      base_query = self.dropout(base_query, deterministic=not train)
      base_vals = self.dropout(base_vals, deterministic=not train)
      if debug:
          out_dict["base_query"] = base_query
          out_dict["base_masks"] = base_masks
      corpus_scores = jax.nn.softmax(
          self.dataset_gate(base_query), axis=-1
      )
      out_dict["corpus_scores"] = corpus_scores
      if no_memory:
          (
              fused_emb,
              fused_mask,
              attn_weights_all_layers,
          ) = self.fusion_encoder(
              fused_input_embs=base_vals,
              fused_mask=base_masks,
              use_dropout=train,
          )  # B×(I+N)×d
      else:
          if use_memory:
              detached_query = jax.lax.stop_gradient(base_query)
              if retrieve_local:
                  (
                      topk_scores,
                      ret_memory,
                      ret_data,
                      local_topk_ids,
                      global_topk_ids,
                      global_corpus_ids,
                  ) = self._dist_mips_local(
                      query=detached_query,
                      corpus_scores=corpus_scores,
                      local_device_id=device_id,
                  )
              else:
                  (
                      topk_scores,
                      ret_memory,
                      ret_data,
                      local_topk_ids,
                      global_topk_ids,
                      global_corpus_ids,
                  ) = self._dist_mips_across(
                      query=detached_query,
                      corpus_scores=corpus_scores,
                      local_device_id=device_id,
                  )
              out_dict["topk_scores"] = topk_scores
  
              # encode the retrieved data
              (
                  retr_keys,
                  retr_vals,
                  retr_masks,
                  _,
                  disentangle_reg,
              ) = self.encode_topk_knowledge(
                  bsz=bsz,
                  retr_images=ret_data["image"],
                  retr_texts=ret_data["text_tokens"],
                  train=train,
                  random_drop_image=False,
                  frozen_base=frozen_base,
              )
  
              global_corpus_scores = layers.batch_index_select(
                  corpus_scores, global_corpus_ids
              )
  
              if debug:
                  out_dict["detached_query"] = detached_query
                  out_dict[
                      "global_corpus_scores"
                  ] = global_corpus_scores
                  out_dict["global_corpus_ids"] = global_corpus_ids
                  out_dict["local_topk_ids"] = local_topk_ids
                  out_dict["global_topk_ids"] = global_topk_ids
                  out_dict["retr_keys"] = retr_keys
                  out_dict["retr_masks"] = retr_masks
                  out_dict["base_vals"] = base_vals
                  out_dict["retr_vals"] = retr_vals
  
              out_dict["retr_data"] = ret_data
              out_dict["base_norm"] = layers.l2_norm(base_vals).mean()
              out_dict["data_norm"] = layers.l2_norm(retr_vals).mean()
              out_dict["vals_norm"] = layers.l2_norm(
                  ret_memory["values"][0]
              ).mean()
              out_dict["gap"] = jnp.abs(
                  1
                  - jnp.divide(
                      out_dict["data_norm"], out_dict["base_norm"]
                  )
              )
  
              if train and retr_texts is not None and use_psudo_retr:
                  logging.info("global keys!!!")
                  (
                      ground_truth_keys,
                      ground_truth_vals,
                      _,
                      _,
                      _,
                  ) = self.encode_knowledge(
                      retr_texts=retr_texts,
                      retr_images=retr_images,
                      bsz=bsz,
                      train=train,
                      random_drop_image=True,
                      frozen_base=frozen_base,
                  )
                  global_keys = jnp.concatenate(
                      jax.lax.all_gather(
                          x=ground_truth_keys, axis_name="batch", axis=0
                      ),
                      axis=0,
                  )
                  logging.info(global_keys.shape)
                  inbatch_sim = jax.lax.dot(
                      base_query, global_keys.transpose()
                  )
                  out_dict["inbatch_sim"] = inbatch_sim
                  if debug:
                      out_dict["global_keys"] = global_keys
                      out_dict["ground_truth_keys"] = ground_truth_keys
                      out_dict["ground_truth_vals"] = ground_truth_vals
                  # replace retrieved knowledge as ground-truth ones for stablization.
                  k = retr_keys.shape[1]
                  ground_truth_keys = jnp.repeat(
                      jnp.expand_dims(ground_truth_keys, axis=1),
                      axis=1,
                      repeats=k,
                  )
                  ground_truth_vals = jnp.repeat(
                      jnp.expand_dims(ground_truth_vals, axis=1),
                      axis=1,
                      repeats=k,
                  )
                  replace_mask = jax.random.bernoulli(
                      self.make_rng("dropout"),
                      p=0.02,
                      shape=(bsz, 1, 1),
                  )
                  keys_mask = jnp.broadcast_to(
                      replace_mask, retr_keys.shape
                  )
                  retr_keys = jax.lax.select(
                      keys_mask, ground_truth_keys, retr_keys
                  )
                  vals_mask = jnp.broadcast_to(
                      jnp.expand_dims(replace_mask, axis=-1),
                      retr_vals.shape,
                  )
                  retr_vals = jax.lax.select(
                      vals_mask, ground_truth_vals, retr_vals
                  )
  
              logging.info("Concat memory and data!!!")
              logging.info(retr_keys.shape)
              logging.info(ret_memory["keys"].shape)
              logging.info(global_corpus_scores.shape)
              # concat retrieved memory (90%) with re-encoded ones (10%)
  
              retr_keys = jnp.concatenate(
                  [retr_keys, ret_memory["keys"]], axis=1
              )
              retr_keys = retr_keys * jnp.expand_dims(
                  global_corpus_scores, axis=-1
              )
              retr_vals = jnp.concatenate(
                  [retr_vals, ret_memory["values"]], axis=1
              )
              retr_masks = jnp.concatenate(
                  [retr_masks, ret_memory["masks"]], axis=1
              )
          elif retr_texts is not None:
              (
                  retr_keys,
                  retr_vals,
                  retr_masks,
                  _,
                  disentangle_reg,
              ) = self.encode_topk_knowledge(
                  bsz=bsz,
                  retr_images=jnp.expand_dims(retr_images, axis=1),
                  retr_texts=jnp.expand_dims(retr_texts, axis=1),
                  train=train,
                  random_drop_image=False,
              )
          else:
              (
                  retr_keys,
                  retr_vals,
                  retr_masks,
                  _,
                  disentangle_reg,
              ) = self.encode_topk_knowledge(
                  bsz=bsz,
                  retr_images=jnp.expand_dims(
                      encoder_input_image, axis=1
                  ),
                  retr_texts=jnp.expand_dims(
                      encoder_input_tokens, axis=1
                  ),
                  train=train,
                  random_drop_image=False,
              )
  
          (
              fused_emb,
              fused_mask,
              retr_scores,
              attn_weights_all_layers,
          ) = self.fuse_topk_knowledge(
              base_query=base_query,
              base_vals=base_vals,
              base_masks=base_masks,
              retr_keys=retr_keys,
              retr_vals=retr_vals,
              retr_masks=retr_masks,
              train=train,
          )  # B×(I+N+M*K)×d
          out_dict["disentangle_reg"] = jnp.mean(disentangle_reg)
          out_dict["retr_scores"] = retr_scores
  
      out_dict["fused_emb"] = fused_emb
      out_dict["fused_mask"] = fused_mask
      logging.info("fused_emb.shape")
      logging.info(fused_emb.shape)
      out_dict["attn_weights_all_layers"] = attn_weights_all_layers
  
      if not only_encode:
          # decode: generate decoding results.
          out_dict["predicted_logits"] = self.t5_decode(
              encoded=fused_emb,
              encoder_input_tokens=fused_mask,
              decoder_input_tokens=decoder_input_tokens,
              decoder_target_tokens=decoder_target_tokens,
              enable_dropout=train,
              decode=decode,
              max_decode_length=max_decode_length,
          )
      return out_dict
  
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 KnowledgeFIDModule obj's underlying member method code with example usage: obj.__call__(decoder_input_tokens, decoder_target_tokens, encoder_input_image, encoder_input_tokens, retr_texts, retr_images, device_id, train, decode, max_decode_length, use_memory, use_psudo_retr, retrieve_local, no_memory, debug, frozen_base, only_encode) and returns: out_dict