code: def embed_3d_patch(
      x, patches, embedding_dim, kernel_init_method, name="embedding"
  ):
      """Embed 3D input patches into tokens."""
  
      assert (
          patches.get("size") is not None
      ), "patches.size must be defined"
      assert len(patches.size) == 3, "patches.size must have 3 elements"
      assert embedding_dim, "embedding_dim must be specified"
  
      fh, fw, ft = patches.size
  
      if kernel_init_method == "central_frame_initializer":
          kernel_initializer = model_utils.central_frame_initializer()
          logging.info(
              "Using central frame initializer for input embedding"
          )
      elif kernel_init_method == "average_frame_initializer":
          kernel_initializer = model_utils.average_frame_initializer()
          logging.info(
              "Using average frame initializer for input embedding"
          )
      else:
          kernel_initializer = default_kernel_init
          logging.info("Using default initializer for input embedding")
  
      x = nn.Conv(
          embedding_dim,
          (ft, fh, fw),
          strides=(ft, fh, fw),
          padding="VALID",
          name=name,
          kernel_init=kernel_initializer,
      )(x)
  
      return x
  
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 3 for calling a global function with example usage: embed_3d_patch(x, patches, embedding_dim, kernel_init_method, name) and returns: x