code: def temporal_encode(
      x,
      temporal_encoding_config,
      patches,
      hidden_size,
      return_1d=True,
      name="embedding",
  ):
      """Encode video for feeding into ViT."""
  
      n, _, in_h, in_w, c = x.shape
  
      if temporal_encoding_config.method == "temporal_sampling":
          n_sampled_frames = temporal_encoding_config.n_sampled_frames
          x = video_utils.sample_frames_uniformly(x, n_sampled_frames)
          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])
  
      elif 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, name
          )
          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.")
  
      assert x.size > 0, (
          "Found zero tokens after temporal encoding. "
          "Perhaps one of the patch sizes is such that "
          "floor(dim_size / patch_size) = 0?"
      )
  
      return x, temporal_dims
  
imports: import functools
  from typing import Any, Optional, Callable, Sequence
  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.common_lib import video_utils
  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.vivit import model_utils
label_desc: Use image node_id 5 for calling a global function with example usage: temporal_encode(x, temporal_encoding_config, patches, hidden_size, return_1d, name) and returns: x,  temporal_dims