class_docstr: FID model (https://arxiv.org/pdf/2007.01282.pdf) with a retrieval module over a knowledge memory.
class_docstr_tok: ["FID","model","(","https",":","\/\/arxiv.org\/pdf\/2007.01282.pdf",")","with","a","retrieval","module","over","a","knowledge","memory","."]
code: def _dist_mips_across(
      self,
      query,
      corpus_scores,
      local_device_id,
      recall_target=0.99,
      exact=False,
  ):
      # must have n_host > retr_k
      logging.info("mips global!!!")
      logging.info(self.local_keys.value.shape)
      logging.info(local_kb.n_data)
      n_local_device = len(self.across_index_groups)
      logging.info(n_local_device)
      global_query = jax.lax.all_gather(
          x=query, axis_name="batch", axis=0, tiled=True
      )
      logging.info(global_query.shape)
      # global_query: (per_bsz * n_local_device * n_hosts) * d
      global_corpus_scores = jax.lax.all_gather(
          x=corpus_scores, axis_name="batch", axis=0, tiled=True
      )
      logging.info(global_corpus_scores.shape)
      # global_corpus_scores: (per_bsz * n_local_device * n_hosts) * n_kb
  
      local_scores = jax.lax.dot(
          global_query, self.local_keys.value.transpose()
      )
      local_k = local_kb.k
      if exact:
          local_topk_scores, local_topk_ids = jax.lax.top_k(
              local_scores, k=local_k
          )
      else:
          local_topk_scores, local_topk_ids = jax.lax.approx_max_k(
              local_scores,
              k=local_k,
              recall_target=recall_target,
              reduction_input_size_override=local_kb.n_data,
              aggregate_to_topk=True,
          )
  
      local_topk_ids_offset = (
          local_topk_ids + local_device_id * local_kb.n_data_per_shard
      )
      logging.info(local_topk_ids.shape)
      # local_topk_ids: (per_bsz * n_local_device * n_hosts) * K
      host_topk_scores = jax.lax.all_gather(
          x=local_topk_scores,
          axis_name="batch",
          axis=1,
          axis_index_groups=self.axis_index_groups,
          tiled=True,
      )
      logging.info(host_topk_scores.shape)
      # host_topk_scores: (per_bsz * n_local_device * n_hosts) * (n_hosts * K)
      host_topk_ids = jax.lax.all_gather(
          x=local_topk_ids_offset,
          axis_name="batch",
          axis=1,
          axis_index_groups=self.axis_index_groups,
          tiled=True,
      )
      # host_topk_ids: (per_bsz * n_local_device * n_hosts) * (n_hosts * K)
  
      host_corpus_scores, host_corpus_ids = self._get_corpus_scores(
          global_corpus_scores, host_topk_ids
      )
      # host_corpus_scores: (per_bsz * n_local_device * n_hosts) * (n_hosts * K)
      host_topk_scores, host_rank_ids = jax.lax.top_k(
          host_topk_scores * host_corpus_scores, k=local_k
      )
      # host_topk_scores: (per_bsz * n_local_device * n_hosts) * K
      host_topk_ids = layers.batch_index_select(
          host_topk_ids, host_rank_ids
      )
      logging.info(host_topk_ids.shape)
      # host_topk_ids: (per_bsz * n_local_device * n_hosts) * K
      host_topk_ids = jnp.reshape(
          host_topk_ids, (-1, n_local_device, local_k)
      )
      host_topk_ids = host_topk_ids[:, local_device_id]
      logging.info(host_topk_ids.shape)
      # host_topk_ids: (per_bsz * n_hosts)
      host_topk_scores = jnp.reshape(
          host_topk_scores, (-1, n_local_device, local_k)
      )
      host_topk_scores = host_topk_scores[:, local_device_id]
      logging.info("host_topk_scores")
      logging.info(host_topk_scores.shape)
  
      ret_memory, ret_data = host_callback.call(
          local_memory.retrieve_top_memory,
          (host_topk_ids),
          result_shape=local_kb.ret_top_specs,
      )
  
      global_topk_scores = jax.lax.all_to_all(
          x=host_topk_scores,
          axis_name="batch",
          split_axis=0,
          concat_axis=1,
          axis_index_groups=self.across_index_groups,
          tiled=True,
      )
      logging.info("global_topk_scores")
      logging.info(global_topk_scores.shape)
      # global_topk_scores: per_bsz * (n_device * k)
      global_topk_scores, global_rank_ids = jax.lax.top_k(
          global_topk_scores, k=self.retr_k
      )
      logging.info(global_topk_scores.shape)
      # global_topk_scores: per_bsz * retr_k
      global_data_ids = global_rank_ids[:, : int(self.data_k)]
      global_memory_ids = global_rank_ids[:, int(self.data_k) :]
  
      def _gather_val(local_ret_vals, top_ids):
          logging.info(local_ret_vals.shape)
          global_ret_vals = jax.lax.all_to_all(
              x=local_ret_vals,
              axis_name="batch",
              split_axis=0,
              concat_axis=1,
              axis_index_groups=self.across_index_groups,
              tiled=True,
          )
          logging.info(global_ret_vals.shape)
          # global_ret_vals: per_bsz * (n_device * k) * dshape
          global_ret_vals = layers.batch_index_select(
              global_ret_vals, top_ids
          )
          logging.info(global_ret_vals.shape)
          # global_ret_vals: per_bsz * retr_k * dshape
          return global_ret_vals
  
      logging.info("_gather_val!!!")
  
      ret_memory = jax.tree_map(
          lambda local_val: _gather_val(local_val, global_memory_ids),
          ret_memory,
      )
  
      ret_data = jax.tree_map(
          lambda local_val: _gather_val(local_val, global_data_ids),
          ret_data,
      )
  
      ret_memory["masks"] = jnp.ones(
          ret_memory["values"].shape[:3]
      ).astype(bool)
  
      for k in ret_memory:
          logging.info(k)
          logging.info(ret_memory[k].shape)
          logging.info(ret_memory[k].dtype)
  
      host_corpus_ids = layers.batch_index_select(
          host_corpus_ids, host_rank_ids
      )
      host_corpus_ids = jnp.reshape(
          host_corpus_ids, (-1, n_local_device, local_k)
      )[:, local_device_id]
      logging.info("corpus_ids")
      logging.info(host_corpus_ids.shape)
      # host_corpus_ids: (per_bsz * n_local_device, k)
  
      global_corpus_ids = jax.lax.all_to_all(
          x=host_corpus_ids,
          axis_name="batch",
          split_axis=0,
          concat_axis=1,
          axis_index_groups=self.across_index_groups,
          tiled=True,
      )
      logging.info(global_corpus_ids.shape)
      # global_corpus_ids: (per_bsz, n_local_device)
      global_corpus_ids = layers.batch_index_select(
          global_corpus_ids, global_rank_ids
      )
      logging.info(global_corpus_ids.shape)
      # global_corpus_ids: (per_bsz, 10)
      return (
          global_topk_scores,
          ret_memory,
          ret_data,
          local_topk_ids,
          global_rank_ids,
          global_corpus_ids,
      )
  
imports: import functools
  from typing import Any, Dict, Mapping, Optional, Tuple, List
  from absl import logging
  import flax.linen
  import jax
  from jax.experimental import host_callback
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.projects.knowledge_visual_language.data import data_utils
  from scenic.projects.knowledge_visual_language.models import constants
  from scenic.projects.knowledge_visual_language.models import fusion_in_decoder_soft
  from scenic.projects.knowledge_visual_language.models import layers
  from scenic.projects.knowledge_visual_language.models import local_memory
  from scenic.projects.knowledge_visual_language.models import losses
  from scenic.projects.knowledge_visual_language.models import metrics
  from t5x import decoding
label_desc: Use image node_id 4 for calling the KnowledgeFIDModule obj's underlying member method code with example usage: obj._dist_mips_across(query, corpus_scores, local_device_id, recall_target, exact) and returns: global_topk_scores,  ret_memory,  ret_data,  local_topk_ids,  global_rank_ids,  global_corpus_ids, global_ret_vals