code: def compute_inners(
      encoded_video: jnp.ndarray,
      encoded_text: jnp.ndarray,
      axis_name: str = "batch",
      return_embeddings: bool = False,
  ):
      """Computes the inner products between the videos and the text prompts.
  
      The videos and texts are first gathered across all devices over the given
      axis name (should agree with the one given to the enclosing pmap).
      We then normalize the tensors along the channel dimension.
  
      Args:
        encoded_video: The encoded videos, shape [batch_size, t, d].
        encoded_text: The encoded text, shape [batch_size, t, n, d]. batch_size is
          the num of samples in the batch, t is the number of test clips, n is the
          number of captions per vid, can include context sentences.
        axis_name: The axis over which the (all-)gather should be done.
        return_embeddings: Whether to return the embeddings or not
  
      Returns:
        A matrix with shape [global_batch_size, global_batch_size * n], s.t.
        position [i, j] holds the inner product of the i'th video with the j%n-th
        prompt of the j//n-th video.
      """
      encoded_video = jax.lax.all_gather(encoded_video, axis_name)
      encoded_text = jax.lax.all_gather(encoded_text, axis_name)
  
      # Merge batch_size and t into a single dimension
      encoded_video = encoded_video.reshape(
          (-1,) + encoded_video.shape[2:]
      )
      encoded_text = encoded_text.reshape(
          (-1,) + encoded_text.shape[2:]
      )
      if return_embeddings:
          return (
              compute_inner_product(encoded_video, encoded_text),
              encoded_video,
              encoded_text,
          )
      else:
          return compute_inner_product(encoded_video, encoded_text)
  
imports: import os
  import re
  from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Union
  from absl import logging
  import flax
  from flax import struct
  from flax.core import frozen_dict
  import flax.linen
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.projects.baselines.clip import model
  from scenic.train_lib import lr_schedules
  from scenic.train_lib import optimizers
  from scenic.train_lib import train_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 14 for calling a global function with example usage: compute_inners(encoded_video, encoded_text, axis_name, return_embeddings) and returns: compute_inner_product,  encoded_video,  encoded_text, compute_inner_product