class_docstr: FID model with a retrieval module over a knowledge memory.
class_docstr_tok: ["FID","model","with","a","retrieval","module","over","a","knowledge","memory","."]
code: def __init__(
      self,
      config: Optional[ml_collections.ConfigDict],
      dataset_meta_data: Dict[str, Any],
      kb_datasets: Dict[str, dataset_utils.Dataset],
  ) -> None:
      self.config = config
      self.dataset_meta_data = dataset_meta_data
      self.retr_k = self.config.model.retr_k
      self.retr_data_ratio = self.config.model.retr_data_ratio
      n_device = jax.device_count()
      self.data_k = int(np.ceil(self.retr_k * self.retr_data_ratio))
      device_per_axis = jax.local_device_count()
      if n_device < device_per_axis:
          self.axis_index_groups = None
          self.across_index_groups = None
      else:
          self.axis_index_groups = np.arange(n_device).reshape(
              [n_device // device_per_axis, device_per_axis]
          )
          self.across_index_groups = self.axis_index_groups.T.tolist()
          self.axis_index_groups = self.axis_index_groups.tolist()
      logging.info("axis_index_groups")
      logging.info(self.axis_index_groups)
      logging.info(self.across_index_groups)
      local_kb.initialize(kb_datasets=kb_datasets)
      self.flax_model = self.build_flax_model()
  
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 to create a new KnowledgeFIDModel object from inherited base classes: base_model with example: obj = KnowledgeFIDModel(config, dataset_meta_data, kb_datasets)