code: def compute_class_text_embeddings(
      train_state: train_utils.TrainState,
      class_texts: jnp.ndarray,  # Shape: (C, T, L)
      mask: Optional[jnp.ndarray] = None,  # Shape: (C, T, L)
      *,
      model: ImageTextModel,
      batch_size: int,
      debug: bool = False,
  ) -> jnp.ndarray:  # Shape: (C, E)
      """Computes the class text embeddings."""
      num_classes = class_texts.shape[0]
      class_texts = class_texts.reshape(-1, class_texts.shape[-1])
      if mask is not None:
          mask = mask.reshape(-1, mask.shape[-1])
  
      def _compute_class_text_embeddings_batch(
          args,  # `(batch, mask)` with shape: (N, L)
      ) -> jnp.ndarray:  # Shape: (N, E)
          batch, mask_ = args
          return model.flax_model.apply(
              variables={
                  "params": train_state.params,
                  **train_state.model_state,
              },
              method=model.flax_model.encode_text,
              text=batch,
              mask=mask_,
              train=False,
              debug=debug,
          )
  
      # We batch to not run OOM.
      encoded_text = split_in_batches(
          _compute_class_text_embeddings_batch, batch_size=batch_size
      )((class_texts, mask))
  
      encoded_text = encoded_text.reshape(
          num_classes, -1, encoded_text.shape[-1]
      ).mean(axis=1)
  
      if axis_name_exists(NUM_DEVICES_AXIS_NAME):
          encoded_text = jax.lax.all_gather(
              encoded_text, NUM_DEVICES_AXIS_NAME
          )
          encoded_text = encoded_text.reshape(
              -1, encoded_text.shape[-1]
          )
  
      return encoded_text
  
imports: from collections.abc import Mapping
  import dataclasses
  import functools
  from typing import Optional
  from absl import logging
  from clu import metric_writers
  from flax import jax_utils
  import flax.core
  import jax
  import jax.numpy
  import ml_collections
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.lang4video import util
  from scenic.projects.lang4video.model.image_text_model import ImageTextModel
  from scenic.projects.lang4video.trainer.train_utils import axis_name_exists
  from scenic.projects.lang4video.trainer.train_utils import compute_mask
  from scenic.projects.lang4video.trainer.train_utils import get_cached_fn
  from scenic.projects.lang4video.trainer.train_utils import get_epoch_steps
  from scenic.projects.lang4video.trainer.train_utils import get_input_spec
  from scenic.projects.lang4video.trainer.train_utils import init_encoder
  from scenic.projects.lang4video.trainer.train_utils import InputSpec
  from scenic.projects.lang4video.trainer.train_utils import is_video_input
  from scenic.projects.lang4video.trainer.train_utils import log_eval_summary
  from scenic.projects.lang4video.trainer.train_utils import NUM_DEVICES_AXIS_NAME
  from scenic.projects.lang4video.trainer.train_utils import pad_array_to_be_divisible
  from scenic.projects.lang4video.trainer.train_utils import partial_with_cache
  from scenic.projects.lang4video.trainer.train_utils import split_in_batches
  from scenic.train_lib import train_utils
  from tqdm.auto import tqdm
label_desc: Use image node_id 2 for calling a global function with example usage: compute_class_text_embeddings(train_state, class_texts, mask) and returns: encoded_text, model