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","."]
code: def setup(self):
      super().setup()
      self.local_keys = self.variable(
          "memory",
          "keys",
          functools.partial(jnp.zeros, dtype=jnp.bfloat16),
          (local_kb.n_data_per_shard, self.key_dim),
      )
      self.local_dataset_idxs = self.variable(
          "memory",
          "idxs",
          functools.partial(jnp.zeros, dtype=jnp.int16),
          (local_kb.n_data_per_shard * local_kb.n_local_device),
      )
      self.dataset_gate = nn.Dense(
          features=local_kb.n_kb_dataset,
          dtype=self.dtype,
          name="dataset_gate",
      )
  
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 1 for calling the KnowledgeFIDModule obj's underlying member method code with example usage: obj.setup() without return types