code: def embed_2d_patch(x, patches, embedding_dim, name="embedding"):
      """Embedding input patches with 2D conv."""
  
      assert patches.get("size") is not None, (
          "patches.size is now the only way" "to define the patches"
      )
      assert embedding_dim, "embedding_dim must be specified"
      fh = patches.size[0]
      fw = patches.size[1]
      x = nn.Conv(
          embedding_dim,
          (fh, fw),
          strides=(fh, fw),
          padding="VALID",
          name=name,
      )(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 2 for calling a global function with example usage: embed_2d_patch(x, patches, embedding_dim, name) and returns: x