class_docstr: 2D Image to Patch Embedding
code: def forward(self, x):
      _, _, H, W = x.shape
      _assert(
          H == self.img_size[0],
          f"Input image height ({H}) doesn't match model ({self.img_size[0]}).",
      )
      _assert(
          W == self.img_size[1],
          f"Input image width ({W}) doesn't match model ({self.img_size[1]}).",
      )
      x = self.proj(
          rearrange(
              x,
              "b c (h p1) (w p2) -> b h w (c p1 p2)",
              p1=self.patch_size[0],
              p2=self.patch_size[1],
          )
      )
      if self.flatten:
          x = rearrange(x, "b h w c -> b (h w) c")
      x = self.norm(x)
      return x
  
file_path: flash-attention/flash_attn/layers/patch_embed.py
imports: from functools import partial
  import torch.nn
  from einops import rearrange
  from torch import _assert
  from torch.nn.modules.utils import _pair
label_desc: Use image node_id 2 for calling the PatchEmbed obj's underlying member method code with example usage: obj.forward(x) and returns: x