class_docstr: Local Knowledge Base stored in CPU.
class_docstr_tok: ["Local","Knowledge","Base","stored","in","CPU","."]
docstr: Load sharded dataset into CPU.
docstr_tok: ["Load","sharded","dataset","into","CPU","."]
code: def initialize(self, kb_datasets):
      """Load sharded dataset into CPU."""
      memory_image = []
      memory_text = []
      memory_idxs = []
      self.n_kb_dataset = len(kb_datasets)
      logging.info("Start Loading sharded dataset into CPU.")
      for idx, dataset_name in enumerate(kb_datasets):
          logging.info(dataset_name)
          dataset = kb_datasets[dataset_name]
          n_iter = int(
              dataset.meta_data["example_per_shard"]
              // dataset.meta_data["batch_size"]
          )
          for _ in tqdm.tqdm(range(n_iter)):
              kb_batch = next(dataset.train_iter)
              memory_image += [np.asarray(kb_batch["image"])]
              memory_text += [np.asarray(kb_batch["knowledge_tokens"])]
              memory_idxs += [
                  idx
                  * np.ones(shape=kb_batch["image"].shape[:2]).astype(
                      "int16"
                  )
              ]
              del kb_batch
      self.memory["image"] = np.concatenate(memory_image, axis=1)
      self.memory["text"] = np.concatenate(memory_text, axis=1)
      self.n_local_device = self.memory["image"].shape[0]
      self.memory_flatten["idxs"] = np.repeat(
          np.reshape(np.concatenate(memory_idxs, axis=1), (1, -1)),
          self.n_local_device,
          axis=0,
      )
      self.n_data_per_shard = self.memory["image"].shape[1]
      self.n_data = self.n_data_per_shard * self.n_local_device
  
      self.specs = {
          "image": dataset.meta_data["image_spec"],
          "text": dataset.meta_data["knowledge_spec"],
      }
  
file_path: scenic/scenic/projects/knowledge_visual_language/models/local_memory.py
imports: import functools
  from absl import logging
  from flax.core.frozen_dict import FrozenDict
  import jax
  import numpy
  import tqdm
label_desc: Use image node_id 3 for calling the KnowledgeBase obj's underlying member method code with example usage: obj.initialize(kb_datasets) without return types