code: def retrieve_memory(args):
      device_id, indexs = args
      return [
          {
              "values": kb.memory["values"][device_id][indexs],
              # 'masks': kb.memory['masks'][device_id][indexs]
          },
          {
              "image": kb.memory["image"][device_id][indexs],
              "text_tokens": kb.memory["text"][device_id][indexs],
          },
      ]
  
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 2 for calling a global function with example usage: retrieve_memory(args) and returns: list