code: def compute_inner_product(encoded_video, encoded_text):
      """Compute inner product between videos and text embeddings."""
      assert len(encoded_video.shape) == 2  # [batch_size, d]
      logging.info("Shape of encoded text is %s", encoded_text.shape)
      assert (
          len(encoded_text.shape) == 3 or len(encoded_text.shape) == 2
      )
  
      # Shape is [batch_size, batch_size, n]
      if len(encoded_text.shape) == 3:
          inners = jnp.einsum(
              "nd,mfd -> nmf", encoded_video, encoded_text
          )
          # Reshaping to [batch_size, batch_size x n]
          return inners.reshape((inners.shape[0], -1))
      return jnp.einsum("nd,md -> nm", 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 15 for calling a global function with example usage: compute_inner_product(encoded_video, encoded_text) and returns: jnp, inners