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 fuse_topk_knowledge(
      self,
      base_query,
      base_vals,
      base_masks,
      retr_keys,
      retr_vals,
      retr_masks,
      train=False,
  ):
      (bsz, k, n_tokens) = retr_vals.shape[:3]
      retr_vals = jnp.reshape(
          retr_vals, (bsz, k * n_tokens, self.ndim)
      )  # B×(M*K)×d
      retr_scores = jnp.einsum("bd,bkd->bk", base_query, retr_keys)
      retr_scores = (
          jax.nn.softmax(self.att_transform(retr_scores), axis=-1) * k
      )
      retr_masks = jnp.reshape(retr_masks, (bsz, k * n_tokens))
      att_mask = [
          jnp.ones([bsz, base_vals.shape[1]]),
          jnp.repeat(retr_scores, repeats=n_tokens, axis=-1),
      ]
      att_mask = jnp.expand_dims(
          jnp.concatenate(att_mask, axis=-1), axis=-1
      )
      (
          fused_query,
          fused_mask,
          attn_weights_all_layers,
      ) = self.fusion_encoder(
          encoder_input_embs=base_vals,
          fused_input_embs=retr_vals,
          encoder_mask=base_masks,
          fused_mask=retr_masks,
          att_mask=att_mask,
          use_dropout=train,
          output=True,
      )  # B×(I+N+M*K)×d
      return (
          fused_query,
          fused_mask,
          retr_scores,
          attn_weights_all_layers,
      )
  
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 8 for calling the FusionInDecoderSoftModule obj's underlying member method code with example usage: obj.fuse_topk_knowledge(base_query, base_vals, base_masks, retr_keys, retr_vals, retr_masks, train) and returns: fused_query,  fused_mask,  retr_scores,  attn_weights_all_layers