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 compress_and_pool_key(self, h, mask):
      window_size = self.n_stride
      pooled_tokens = nn.avg_pool(
          h[:, self.n_compressed_tokens :, :],
          window_shape=(window_size,),
          strides=(self.n_stride,),
      )
      pooled_tokens = jnp.concatenate(
          (h[:, : self.n_compressed_tokens, :], pooled_tokens), axis=1
      )
      pooled_mask = jnp.squeeze(
          -nn.max_pool(
              jnp.expand_dims(
                  -mask[:, self.n_compressed_tokens :], axis=-1
              ),
              window_shape=(window_size,),
              strides=(self.n_stride,),
          )
      )
      pooled_mask = jnp.concatenate(
          (mask[:, : self.n_compressed_tokens], pooled_mask), axis=1
      )
      # Total: 10 + 512 / n_stride = 42 tokens
      return pooled_tokens, pooled_mask
  
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 2 for calling the FusionInDecoderSoftModule obj's underlying member method code with example usage: obj.compress_and_pool_key(h, mask) and returns: pooled_tokens,  pooled_mask