code: def embed_2d_patch(x, patches, embedding_dim):
      """Standard ViT method of embedding input patches."""
  
      n, h, w, c = x.shape
  
      assert patches.get("size") is not None, (
          "patches.size is now the only way" "to define the patches"
      )
  
      fh, fw = patches.size
      gh, gw = h // fh, w // fw
  
      if embedding_dim:
          x = nn.Conv(
              embedding_dim,
              (fh, fw),
              strides=(fh, fw),
              padding="VALID",
              name="embedding",
          )(x)
      else:
          # This path often results in excessive padding: b/165788633
          x = jnp.reshape(x, [n, gh, fh, gw, fw, c])
          x = jnp.transpose(x, [0, 1, 3, 2, 4, 5])
          x = jnp.reshape(x, [n, gh, gw, -1])
  
      return x
  
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 3 for calling a global function with example usage: embed_2d_patch(x, patches, embedding_dim) and returns: x