code: def static_encode_knowledge(
      knowledge_batch, train_state, *, flax_model
  ):
      """Function to encode KB knowledge.
  
      Args:
        knowledge_batch: A single batch of data. The buffer of this argument can be
          donated to the computation.
        train_state: The state of training including the current global_step,
          model_state, rng, and optimizer. The buffer of this argument can be
          donated to the computation.
        flax_model: A Flax model.
  
      Returns:
        Key (single embedding), Val (compressed list of embedding) and Mask.
      """
      variables = {
          "params": train_state.params,
          **train_state.model_state,
      }
      retr_tokens, retr_images = (
          knowledge_batch["knowledge_tokens"],
          knowledge_batch["image"],
      )
      batch_size = retr_images.shape[0]
      (
          keys_head,
          compressed_val,
          compressed_mask,
          _,
          _,
      ) = flax_model.apply(
          variables,
          retr_texts=retr_tokens,
          retr_images=retr_images,
          bsz=batch_size,
          train=False,
          random_drop_image=False,
          method=flax_model.encode_knowledge,
      )
      return keys_head, compressed_val, compressed_mask
  
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 for calling a global function with example usage: static_encode_knowledge(knowledge_batch, train_state) and returns: keys_head,  compressed_val,  compressed_mask