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","."]
code: def encode_topk_knowledge(
      self,
      bsz,
      retr_texts,
      retr_images=None,
      train=False,
      random_drop_image=False,
      frozen_base=True,
  ):
      k, l = retr_texts.shape[1], retr_texts.shape[2]
      retr_texts = jnp.reshape(retr_texts, (bsz * k, l))
      if retr_images is not None:
          image_shape = (bsz * k,) + retr_images.shape[2:]
          retr_images = jnp.reshape(retr_images, image_shape)
      (
          retr_keys,
          compressed_val,
          compressed_mask,
          retr_img_emb,
          disentangle_reg,
      ) = self.encode_knowledge(
          retr_texts,
          retr_images,
          bsz=bsz * k,
          train=train,
          random_drop_image=random_drop_image,
          frozen_base=frozen_base,
      )
      n_tokens = compressed_val.shape[1]
      retr_keys = jnp.reshape(retr_keys, (bsz, k, self.key_dim))
      compressed_val = jnp.reshape(
          compressed_val, (bsz, k, n_tokens, self.ndim)
      )  # B×K×M×d
      compressed_mask = jnp.reshape(compressed_mask, (bsz, k, n_tokens))
      return (
          retr_keys,
          compressed_val,
          compressed_mask,
          retr_img_emb,
          disentangle_reg,
      )
  
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 6 for calling the FusionInDecoderSoftModule obj's underlying member method code with example usage: obj.encode_topk_knowledge(bsz, retr_texts, retr_images, train, random_drop_image, frozen_base) and returns: retr_keys,  compressed_val,  compressed_mask,  retr_img_emb,  disentangle_reg