    def __call__(
        self,
        input_ids,  # [B, S] or [B * S]
        token_type_ids,  # [B, S] or [B * S]
        position_ids,  # [B, S] or [B * S]
        word_embeddings,  # [vocab_size, hidden_size]
        token_type_embeddings,  # [type_vocab_size, hidden_size]
        position_embeddings,  # [max_position_embeddings, hidden_size]
        gamma,  # [hidden_size]
        beta,  # [hidden_size]
        eps=1e-5,
    ) -> Tensor:
        """Compute
        embedding = layernorm(word_embedding + token_type_embedding + position_embedding).

        Parameters
        ----------
        input_ids, token_type_ids, position_ids must have the same sizes, either 2D or 1D.

        Returns
        -------
        The computed embedding.
        """

        # dtype check
        dtype_input_ids = input_ids._attrs["dtype"]
        dtype_token_type_ids = token_type_ids._attrs["dtype"]
        dtype_position_ids = position_ids._attrs["dtype"]
        assert (
            dtype_input_ids == dtype_token_type_ids
            and dtype_input_ids == dtype_position_ids
        ), "dtype of input_ids, token_type_ids, and position_ids must be the same"

        dtype_word_embeddings = word_embeddings._attrs["dtype"]
        dtype_token_type_embeddings = token_type_embeddings._attrs["dtype"]
        dtype_position_embeddings = position_embeddings._attrs["dtype"]
        assert (
            dtype_word_embeddings == dtype_token_type_embeddings
            and dtype_word_embeddings == dtype_position_embeddings
        ), "dtype of word_embeddings, token_type_embeddings, position_embeddings must be the same"

        assert dtype_input_ids in [
            "int",
            "int32",
            "int64",
        ], f"Expected dtype int/int32/int64 for index, got dtype {dtype_input_ids}"

        assert (
            dtype_word_embeddings
            in [
                "float16",
                "float32",
            ]
        ), f"Expected dtype float16/float32 for embeddings, got dtype {dtype_word_embeddings}"

        # expecting all three ids to have the same shapes
        assert shape_utils.is_same_shape(input_ids.shape(), token_type_ids.shape()), (
            f"Expecting input_ids and token_type_ids to have the same shapes, but got "
            f"input_ids.shape(): {input_ids.shape()}, token_type_ids.shape(): {token_type_ids.shape()}"
        )
        assert shape_utils.is_same_shape(input_ids.shape(), position_ids.shape()), (
            f"Expecting input_ids and position_ids to have the same shapes, but got "
            f"input_ids.shape(): {input_ids.shape()}, position_ids.shape(): {position_ids.shape()}"
        )

        # expecting the last dim of all three embedding tables to be the same
        dim = word_embeddings._size(-1)
        assert isinstance(dim, IntImm), f"Embedding dim {dim} must be static."
        dim_value = dim.value()
        assert dim_value % 8 == 0, f"Embedding dim {dim} must be multiple of 8."
        assert dim == token_type_embeddings._size(-1), (
            f"Expecting the last dim of word_embeddings and token_type_embeddings to be the same, "
            f"but got {word_embeddings._size(-1)} and {token_type_embeddings._size(-1)}"
        )
        assert dim == position_embeddings._size(-1), (
            f"Expecting the last dim of word_embeddings and position_embeddings to be the same, "
            f"but got {word_embeddings._size(-1)} and {position_embeddings._size(-1)}"
        )

        self._attrs["eps"] = eps

        self._attrs["inputs"] = [
            input_ids,
            token_type_ids,
            position_ids,
            word_embeddings,
            token_type_embeddings,
            position_embeddings,
            gamma,
            beta,
        ]

        self._set_depth()

        output_shape = self._infer_shapes(input_ids, word_embeddings)
        output = Tensor(
            output_shape,
            src_ops={self},
            dtype=word_embeddings._attrs["dtype"],
        )
        self._attrs["outputs"] = [output]
        return output
Compute
embedding = layernorm(word_embedding + token_type_embedding + position_embedding).

Parameters
----------
input_ids, token_type_ids, position_ids must have the same sizes, either 2D or 1D.

Returns
-------
The computed embedding.
python/aitemplate/compiler/ops/embedding/bert_embeddings.py