code: def temporal_encode(
      x,
      modality,
      temporal_encoding_config,
      patches,
      hidden_size,
      return_1d=True,
  ):
      """Encode video for feeding into ViT."""
      if modality == "spectrogram":
          # Spectrogram is treated as a big num_time_bins by num_mel_bins image.
          x = embed_2d_patch(
              x, patches, hidden_size, "embedding_spectrogram"
          )
          temporal_dims = 1
          if return_1d:
              n, h, w, c = x.shape
              x = jnp.reshape(x, [n, h * w, c])
      elif modality == "rgb":
          if temporal_encoding_config.method == "temporal_sampling":
              n, num_frames, in_h, in_w, c = x.shape
              n_sampled_frames = (
                  temporal_encoding_config.n_sampled_frames
              )
              if n_sampled_frames < num_frames:
                  t_start_idx = num_frames / (n_sampled_frames + 1)
                  t_step = t_start_idx
              else:
                  t_start_idx = 0
                  t_step = 1
              t_end_idx = num_frames
              temporal_indices = jnp.arange(
                  t_start_idx, t_end_idx, t_step
              )
              temporal_indices = jnp.round(temporal_indices).astype(
                  jnp.int32
              )
              temporal_indices = jnp.minimum(
                  temporal_indices, num_frames - 1
              )
  
              x = x[:, temporal_indices]  # [n, t_s, in_h, in_w, c]
              t_s = x.shape[1]
              x = jnp.reshape(x, [n, t_s * in_h, in_w, c])
              x = embed_2d_patch(x, patches, hidden_size)
              temporal_dims = t_s
              if return_1d:
                  n, th, w, c = x.shape
                  x = jnp.reshape(x, [n, th * w, c])
              else:
                  n, th, w, c = x.shape
                  x = jnp.reshape(x, [n, t_s, -1, w, c])
          if temporal_encoding_config.method == "3d_conv":
              kernel_init_method = temporal_encoding_config.get(
                  "kernel_init_method", None
              )
  
              x = embed_3d_patch(
                  x, patches, hidden_size, kernel_init_method
              )
              temporal_dims = x.shape[1]
              if return_1d:
                  n, t, h, w, c = x.shape
                  x = jnp.reshape(x, [n, t * h * w, c])
      else:
          raise AssertionError("Unknown temporal encoding method.")
      return x, temporal_dims
  
imports: import functools
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple
  from absl import logging
  import flax.linen
  from flax.linen.linear import default_kernel_init
  from immutabledict import immutabledict
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.mbt import model_utils
label_desc: Use image node_id 4 for calling a global function with example usage: temporal_encode(x, modality, temporal_encoding_config, patches, hidden_size, return_1d) and returns: x,  temporal_dims