class_docstr: Local Knowledge Base stored in CPU.
class_docstr_tok: ["Local","Knowledge","Base","stored","in","CPU","."]
code: def set_encode_fn(self, flax_model):
      encode_fn = functools.partial(
          static_encode_knowledge, flax_model=flax_model
      )
      self.encode_knowledge_pmap = jax.pmap(
          encode_fn,
          axis_name="batch",
          donate_argnums=(0, 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 2 for calling the KnowledgeBase obj's underlying member method code with example usage: obj.set_encode_fn(flax_model) without return types