class_docstr: Local Knowledge Base stored in CPU.
class_docstr_tok: ["Local","Knowledge","Base","stored","in","CPU","."]
docstr: Function to update stale embedding as memory.
  
  Args:
    pmap_train_state: pmaped train state.
    bsz: Global batch size.
    retr_k: number of returned data for retrieval.
    data_k: number of returned data for ranking.
    axis_index_groups: axis groups to gather data.
  
  Returns:
    updated train_state
docstr_tok: ["Function","to","update","stale","embedding","as","memory",".","Args",":","pmap_train_state",":","pmaped","train","state",".","bsz",":","Global","batch","size",".","retr_k",":","number","of","returned","data","for","retrieval",".","data_k",":","number","of","returned","data","for","ranking",".","axis_index_groups",":","axis","groups","to","gather","data",".","Returns",":","updated","train_state"]
code: def update_memory(
      self, pmap_train_state, bsz, retr_k, data_k, axis_index_groups
  ):
      """Function to update stale embedding as memory.
  
      Args:
        pmap_train_state: pmaped train state.
        bsz: Global batch size.
        retr_k: number of returned data for retrieval.
        data_k: number of returned data for ranking.
        axis_index_groups: axis groups to gather data.
  
      Returns:
        updated train_state
      """
      per_bsz = bsz // jax.device_count()
      if axis_index_groups is None:
          per_shard_bsz = bsz
      else:
          per_shard_bsz = bsz // len(axis_index_groups[0])
      logging.info("update memory!!!")
      logging.info(per_bsz)
      memory_key = []
      memory_val = []
      # memory_mask = []
      eval_per_bsz = per_bsz * 4
      for idx in range(
          int(np.ceil(self.n_data_per_shard / eval_per_bsz))
      ):
          kb_batch = {
              "knowledge_tokens": self.memory["text"][
                  :, idx * eval_per_bsz : (idx + 1) * eval_per_bsz
              ],
              "image": self.memory["image"][
                  :, idx * eval_per_bsz : (idx + 1) * eval_per_bsz
              ],
          }
          keys_head, compressed_val, _ = self.encode_knowledge_pmap(
              kb_batch, pmap_train_state
          )
          memory_key += [np.asarray(keys_head)]
          memory_val += [np.asarray(compressed_val)]
          # memory_mask += [np.asarray(mask)]
          del kb_batch
  
      for kw in ["keys", "values"]:
          if kw in self.memory:
              del self.memory[kw]
          if kw in self.memory_flatten:
              del self.memory_flatten[kw]
  
      self.memory["keys"] = np.concatenate(memory_key, axis=1)
      self.memory["values"] = np.concatenate(memory_val, axis=1)
      # self.memory['masks'] = np.concatenate(memory_mask, axis=1)
  
      for kw in ["keys", "values", "image", "text"]:
          mem = self.memory[kw]
          self.memory_flatten[kw] = mem.reshape(
              (mem.shape[0] * mem.shape[1],) + mem.shape[2:]
          )
  
      self.specs["keys"] = (keys_head.shape[2:], keys_head.dtype.name)
      self.specs["values"] = (
          compressed_val.shape[2:],
          compressed_val.dtype.name,
      )
      # self.specs['masks'] = (mask.shape[2:], mask.dtype.name)
  
      self.local_ret_specs = [
          {
              "keys": jax.ShapeDtypeStruct(
                  shape=(per_bsz, retr_k - data_k)
                  + self.specs["keys"][0],
                  dtype=self.specs["values"][1],
              ),
              "values": jax.ShapeDtypeStruct(
                  shape=(per_bsz, retr_k - data_k)
                  + self.specs["values"][0],
                  dtype=self.specs["values"][1],
              ),
          },
          {
              "image": jax.ShapeDtypeStruct(
                  shape=(per_bsz, data_k) + self.specs["image"][0],
                  dtype=self.specs["image"][1],
              ),
              "text_tokens": jax.ShapeDtypeStruct(
                  shape=(per_bsz, data_k) + self.specs["text"][0],
                  dtype=self.specs["text"][1],
              ),
          },
      ]
      self.k = int(np.ceil(retr_k / len(axis_index_groups)) + 1)
      self.ret_top_specs = [
          {
              "keys": jax.ShapeDtypeStruct(
                  shape=(per_shard_bsz, self.k) + self.specs["keys"][0],
                  dtype=self.specs["values"][1],
              ),
              "values": jax.ShapeDtypeStruct(
                  shape=(per_shard_bsz, self.k)
                  + self.specs["values"][0],
                  dtype=self.specs["values"][1],
              ),
          },
          {
              "image": jax.ShapeDtypeStruct(
                  shape=(per_shard_bsz, self.k)
                  + self.specs["image"][0],
                  dtype=self.specs["image"][1],
              ),
              "text_tokens": jax.ShapeDtypeStruct(
                  shape=(per_shard_bsz, self.k) + self.specs["text"][0],
                  dtype=self.specs["text"][1],
              ),
          },
      ]
  
      logging.info(self.local_ret_specs)
      logging.info(self.ret_specs)
      logging.info(self.memory["keys"].shape)
      new_model_state = pmap_train_state.model_state.unfreeze()
      if "keys" in new_model_state["memory"]:
          del new_model_state["memory"]["keys"]
          del new_model_state["memory"]["idxs"]
      new_model_state["memory"]["keys"] = self.memory["keys"]
      new_model_state["memory"]["idxs"] = self.memory_flatten["idxs"]
      pmap_train_state = pmap_train_state.replace(
          model_state=FrozenDict(new_model_state)
      )
      logging.info("finish update memory!!!")
      return pmap_train_state
  
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 4 for calling the KnowledgeBase obj's underlying member method code with example usage: obj.update_memory(pmap_train_state, bsz, retr_k, data_k, axis_index_groups) and returns: pmap_train_state