code: def get_logits_fn(
      train_state: train_utils.TrainState,
      batch: Batch,
      dims: Optional[List[int]],
      *,
      flax_model: nn.Module,
      gather_to_host: bool = True,
  ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
      """Gets the logits from an input abtch.
  
      Args:
        train_state: Loaded train state from the ckpt.
        batch: Input batch.
        dims: Matvit nesting dimensions, a list of size num_layers.
        flax_model: The trained mode.
        gather_to_host: Whether to gather all outputs to host device.
  
      Returns:
        Logits for this batch.
      """
      variables = {
          "params": train_state.params,
          **train_state.model_state,
      }
      logits = flax_model.apply(
          variables,
          batch["inputs"],
          train=False,
          debug=False,
          matvit_mask_dims=dims,
      )
      if gather_to_host:
          logits = jax.lax.all_gather(logits, "batch")
          batch = jax.lax.all_gather(batch, "batch")
      return logits, batch["label"], batch["batch_mask"]
  
file_path: scenic/scenic/projects/matvit/classification_eval_main.py
imports: import functools
  from typing import Callable, Dict, List, Optional, Tuple
  from absl import app
  from absl import flags
  from absl import logging
  from flax import jax_utils
  import flax.linen
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.projects.matvit import matvit
  from scenic.train_lib import train_utils
label_desc: Use image node_id 2 for calling a global function with example usage: get_logits_fn(train_state, batch, dims) and returns: logits,  batch,  batch