code: def embed_2d_patch(
      x, patches, embedding_dim, return_1d=True, 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)
  
      if return_1d:
          batch_size = x.shape[0]
          x = jnp.reshape(x, [batch_size, -1, embedding_dim])
      return x
  
file_path: scenic/scenic/projects/av_mae/model_utils.py
imports: from typing import Callable, Iterable, Optional, Sequence
  import flax.linen
  import jax
  import jax.numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.projects.baselines import vit
label_desc: Use image node_id 5 for calling a global function with example usage: embed_2d_patch(x, patches, embedding_dim, return_1d, name) and returns: x