class_docstr: Modification of FID (https://arxiv.org/pdf/2007.01282.pdf) model.
  
  Take continous embedding of retrieved document at middle fusion layer
  instead of whole sequence at input.
class_docstr_tok: ["Modification","of","FID","(","https",":","\/\/arxiv.org\/pdf\/2007.01282.pdf",")","model",".","Take","continous","embedding","of","retrieved","document","at","middle","fusion","layer","instead","of","whole","sequence","at","input","."]
docstr: Conduct supervised retrieval-augmented training with given retrieved documents.
  
  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.
    train: whether using train mode.
    decode: whether in decode mode.
    fuse_retrieval: whether use input retrieval docs.
    max_decode_length: maximum decode token length.
    debug: whether use debug mode.
    in_batch_neg: whether use in-batch contastive learning.
    frozen_base: whether froze the whole encoder.
    **args: other possible arguments.
  
  Returns:
    output dictionary containing final and intermediate results.
docstr_tok: ["Conduct","supervised","retrieval-augmented","training","with","given","retrieved","documents",".","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",".","train",":","whether","using","train","mode",".","decode",":","whether","in","decode","mode",".","fuse_retrieval",":","whether","use","input","retrieval","docs",".","max_decode_length",":","maximum","decode","token","length",".","debug",":","whether","use","debug","mode",".","in_batch_neg",":","whether","use","in-batch","contastive","learning",".","frozen_base",":","whether","froze","the","whole","encoder",".","*","*","args",":","other","possible","arguments",".","Returns",":","output","dictionary","containing","final","and","intermediate","results","."]
code: def __call__(
      self,
      decoder_input_tokens,  # B×O
      decoder_target_tokens,  # B×O
      encoder_input_image=None,  # B×W×H×3
      encoder_input_tokens=None,  # B×I
      retr_texts=None,  # B×K×L
      retr_images=None,  # B×K×W×H×3
      train=False,
      decode=False,
      fuse_retrieval=True,
      max_decode_length=None,
      debug: bool = False,
      in_batch_neg: bool = False,
      frozen_base=True,
      **args,
  ):
      """Conduct supervised retrieval-augmented training with given retrieved documents.
  
      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.
        train: whether using train mode.
        decode: whether in decode mode.
        fuse_retrieval: whether use input retrieval docs.
        max_decode_length: maximum decode token length.
        debug: whether use debug mode.
        in_batch_neg: whether use in-batch contastive learning.
        frozen_base: whether froze the whole encoder.
        **args: other possible arguments.
  
      Returns:
        output dictionary containing final and intermediate results.
      """
      bsz = decoder_input_tokens.shape[0]
      base_vals, base_masks, query_img_emb = self.get_base_encoded(
          bsz=bsz,
          image=encoder_input_image,
          text_tokens=encoder_input_tokens,
          train=train,
          frozen_base=frozen_base,
      )  # B×N×d, B×I×d
      out_dict = {
          "query_img_emb": query_img_emb,
          "text_query": base_vals[0],
          "image_query": base_vals[1],
      }
      base_vals = jnp.concatenate(base_vals, axis=1)  # B×(I+N)×d
      if retr_texts is not None:
          (
              retr_keys,
              retr_vals,
              retr_masks,
              retr_img_emb,
              disentangle_reg,
          ) = self.encode_topk_knowledge(
              bsz=bsz,
              retr_images=retr_images,
              retr_texts=retr_texts,
              train=train,
              random_drop_image=True,
          )
          base_query = self.query_head(
              encoded_emb=base_vals,
              encoder_mask=base_masks,
              use_dropout=train,
          )  # B×(I+N)×d -> B×d
          out_dict["disentangle_reg"] = disentangle_reg
          out_dict["retr_img_emb"] = retr_img_emb
          out_dict["base_query"] = base_query
          out_dict["retr_keys"] = retr_keys
          out_dict["retr_vals"] = retr_vals
  
      if fuse_retrieval and retr_texts is not None:
          # fuse top-k retrieved knowledge (or no fusion)
          if in_batch_neg and retr_vals.shape[1] == 1:
              # retr_vals: B×1×M×d -> B×2×M×d, retr_keys: B×1×d -> B×2×M×d
              retr_vals = jnp.concatenate(
                  (retr_vals, jnp.roll(retr_vals, shift=1, axis=0)),
                  axis=1,
              )
              retr_keys = jnp.concatenate(
                  (retr_keys, jnp.roll(retr_keys, shift=1, axis=0)),
                  axis=1,
              )
              retr_masks = jnp.concatenate(
                  (retr_masks, jnp.roll(retr_masks, shift=1, axis=0)),
                  axis=1,
              )
  
          (
              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["retr_scores"] = retr_scores
      else:
          # only fuse input image and text
          (
              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
      # generate decoding results.
      out_dict["attn_weights_all_layers"] = attn_weights_all_layers
      out_dict["predicted_logits"] = self.out_decoder(
          encoded=fused_emb,
          decoder_input_tokens=decoder_input_tokens,
          encoder_input_tokens=fused_mask,
          decoder_target_tokens=decoder_target_tokens,
          enable_dropout=train,
          decode=decode,
          max_decode_length=max_decode_length,
          encoder_segment_ids=None,
          decoder_segment_ids=None,
      )
      return out_dict
  
imports: from typing import Any, Dict, Optional
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import base_model
  from scenic.projects.knowledge_visual_language.models import constants
  from scenic.projects.knowledge_visual_language.models import layers
  from scenic.projects.knowledge_visual_language.models import losses
  from scenic.projects.knowledge_visual_language.models import metrics
  from scenic.projects.knowledge_visual_language.models import vit
  from scenic.projects.t5 import layers
  from scenic.projects.t5 import model
label_desc: Use image node_id 9 for calling the FusionInDecoderSoftModule 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, train, decode, fuse_retrieval, max_decode_length, debug, in_batch_neg, frozen_base) and returns: out_dict