Method __init__

tensorrt_llm/layers/normalization.py:177–212

(self,
                 embedding_dim: int,
                 num_embeddings: Optional[int] = None,
                 norm_type: str = "layer_norm",
                 bias: bool = True,
                 mapping=Mapping(),
                 dtype=None)

class AdaLayerNormZero(Module):

    def __init__(self,
                 embedding_dim: int,
                 num_embeddings: Optional[int] = None,
                 norm_type: str = "layer_norm",
                 bias: bool = True,
                 mapping=Mapping(),
                 dtype=None):
        super().__init__()
        if num_embeddings is not None:
            self.emb = CombinedTimestepLabelEmbeddings(num_embeddings,
                                                       embedding_dim,
                                                       dtype=dtype)
        else:
            self.emb = None

        self.silu = ACT2FN['silu']
        self.linear = Linear(embedding_dim,
                             6 * embedding_dim,
                             bias=bias,
                             tp_group=mapping.tp_group,
                             tp_size=mapping.tp_size,
                             dtype=dtype)
        if norm_type == "layer_norm":
            self.norm = LayerNorm(embedding_dim,
                                  elementwise_affine=False,
                                  eps=1e-6,
                                  dtype=dtype)
        elif norm_type == "fp32_layer_norm":
            self.norm = LayerNorm(embedding_dim,
                                  elementwise_affine=False,
                                  bias=False,
                                  dtype=dtype)
        else:
            raise ValueError(
                f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'."
            )

    def forward(self,
                x: Tensor,