code: def compute_similarity_scores(
      train_state: train_utils.TrainState,
      iterator,
      eval_step_fn,
      eval_steps,
      workdir,
      lead_host,
  ):
      """Computes similarity scores and dump them directly instead of metrics."""
      # Sync model state across replicas.
      train_state = train_utils.sync_model_state_across_replicas(
          train_state
      )
      (
          all_logits,
          all_keys,
          all_labels,
          all_batch_masks,
          all_intermediate,
      ) = ([], [], [], [], [])
      # Do this to ensure we definitely cover the full test set
      eval_steps = int(np.ceil(1.3 * eval_steps))
      logging.info("Number of eval steps is %s", eval_steps)
      for step in range(eval_steps):
          with jax.profiler.StepTraceAnnotation("eval", step_num=step):
              eval_batch = next(iterator)
              assert "key" in eval_batch, "Keys must be added to batch"
              keys = eval_batch["key"]
              labels = eval_batch["label"]
              batch_masks = eval_batch["batch_mask"]
              del eval_batch["key"]
              del eval_batch["label"]
  
              logits, intermediate = eval_step_fn(
                  train_state, eval_batch
              )
              (
                  gathered_logits,
                  gathered_keys,
                  gathered_labels,
                  gathered_batch_masks,
                  gathered_intermediate,
              ) = all_gather_and_unreplicate(
                  (logits, keys, labels, batch_masks, intermediate)
              )
              all_logits.append(np.concatenate(gathered_logits, axis=0))
              all_intermediate.append(
                  np.concatenate(gathered_intermediate, axis=0)
              )
              all_labels.append(np.concatenate(gathered_labels, axis=0))
              all_keys.append(
                  tf.strings.unicode_encode(
                      np.concatenate(gathered_keys, axis=0), "UTF-8"
                  )
              )
              all_batch_masks.append(
                  np.concatenate(gathered_batch_masks, axis=0)
              )
  
      logging.info("all_scores.shape: %s", str(len(all_keys)))
  
      timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
      fname_logits = os.path.join(
          workdir, f"logits_offline_eval_{timestamp}"
      )
      fname_intermediate = os.path.join(
          workdir, f"intermediate_offline_eval_{timestamp}"
      )
      fname_labels = os.path.join(
          workdir, f"labels_offline_eval_{timestamp}"
      )
      fname_keys = os.path.join(
          workdir, f"keys_offline_eval_{timestamp}"
      )
      fname_masks = os.path.join(
          workdir, f"masks_offline_eval_{timestamp}"
      )
      if lead_host:
          logging.info("Logging results to %s", fname_logits)
          log_to_cns(
              predictions=np.concatenate(all_logits, axis=0),
              filename_prefix=fname_logits,
          )
          log_to_cns(
              predictions=np.concatenate(all_intermediate, axis=0),
              filename_prefix=fname_intermediate,
          )
          log_to_cns(
              predictions=np.concatenate(all_keys, axis=0),
              filename_prefix=fname_keys,
          )
          log_to_cns(
              predictions=np.concatenate(all_labels, axis=0),
              filename_prefix=fname_labels,
          )
          log_to_cns(
              predictions=np.concatenate(all_batch_masks, axis=0),
              filename_prefix=fname_masks,
          )
  
imports: import datetime
  import functools
  import os
  import pickle
  from typing import Any, Callable, Optional, Type
  from absl import logging
  from clu import metric_writers
  from flax import jax_utils
  import flax.linen
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.common_lib import debug_utils
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.projects.svvit import classification_trainer
  from scenic.projects.svvit import metrics
  from scenic.train_lib import optimizers
  from scenic.train_lib import pretrain_utils
  from scenic.train_lib import train_utils
  import tensorflow
label_desc: Use image node_id 3 for calling a global function with example usage: compute_similarity_scores(train_state, iterator, eval_step_fn, eval_steps, workdir, lead_host) without return types