class_docstr: Local Knowledge Base stored in CPU.
class_docstr_tok: ["Local","Knowledge","Base","stored","in","CPU","."]
code: def __init__(self):
      self.memory = {}
      self.memory_flatten = {}
      self.specs = {}
      self.n_data_per_shard = -1
      self.n_data = -1
      self.ret_specs = None
      self.n_kb_dataset = 1
      self.n_local_device = 1
      self.k = 1
  
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 1 to create a new KnowledgeBase object  with example: obj = KnowledgeBase()