    def __init__(
        self,
        embed_dim: int,
        patch_embed_shape: Tuple[int, int, int],
        sep_pos_embed: bool = False,
        has_cls: bool = True,
        dtype: str = "float16",
    ) -> None:
        """
        Args:
            embed_dim (int): Embedding dimension for input sequence.
            patch_embed_shape (Tuple): The number of patches in each dimension
                (T, H, W) after patch embedding.
            sep_pos_embed (bool): If set to true, one positional encoding is used for
                spatial patches and another positional encoding is used for temporal
                sequence. Otherwise, only one positional encoding is used for all the
                patches.
            has_cls (bool): If set to true, a cls token is added in the beginning of each
                input sequence.
        """
        super().__init__()
        assert (
            len(patch_embed_shape) == 3
        ), "Patch_embed_shape should be in the form of (T, H, W)."
        self.cls_embed_on = has_cls
        self.sep_pos_embed = sep_pos_embed
        self._patch_embed_shape = tuple(patch_embed_shape)
        self.num_spatial_patch = patch_embed_shape[1] * patch_embed_shape[2]
        self.num_temporal_patch = patch_embed_shape[0]

        if self.cls_embed_on:
            self.cls_token = Parameter(shape=[1, 1, embed_dim], dtype=dtype)
            num_patches = self.num_spatial_patch * self.num_temporal_patch + 1
        else:
            self.cls_token = Parameter(shape=[], value=0, dtype=dtype)
            num_patches = self.num_spatial_patch * self.num_temporal_patch

        if self.sep_pos_embed:
            self.pos_embed_spatial = Parameter(
                shape=[1, self.num_spatial_patch, embed_dim],
                dtype=dtype,
            )
            self.pos_embed_temporal = Parameter(
                shape=[1, self.num_temporal_patch, embed_dim],
                dtype=dtype,
            )
            if self.cls_embed_on:
                self.pos_embed_class = Parameter(shape=[1, 1, embed_dim], dtype=dtype)
            else:
                self.pos_embed_class = Parameter(shape=[], dtype=dtype)
            self.pos_embed = Parameter(shape=[], dtype=dtype)

        else:
            self.pos_embed = Parameter(shape=[1, num_patches, embed_dim], dtype=dtype)
            # Placeholders for torchscriptability, won't be used
            self.pos_embed_spatial = Parameter(shape=[], dtype=dtype)
            self.pos_embed_temporal = Parameter(shape=[], dtype=dtype)
            self.pos_embed_class = Parameter(shape=[], dtype=dtype)
Args:
embed_dim (int): Embedding dimension for input sequence.
patch_embed_shape (Tuple): The number of patches in each dimension
(T, H, W) after patch embedding.
sep_pos_embed (bool): If set to true, one positional encoding is used for
spatial patches and another positional encoding is used for temporal
sequence. Otherwise, only one positional encoding is used for all the
patches.
has_cls (bool): If set to true, a cls token is added in the beginning of each
input sequence.
python/aitemplate/frontend/nn/positional_encoding.py