Method __init__

tensorrt_llm/layers/normalization.py:308–330

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

class SD35AdaLayerNormZeroX(Module):

    def __init__(self,
                 embedding_dim: int,
                 norm_type: str = "layer_norm",
                 bias: bool = True,
                 mapping=Mapping(),
                 dtype=None):
        super().__init__()
        self.silu = ACT2FN['silu']
        self.linear = Linear(embedding_dim,
                             9 * 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)
        else:
            raise ValueError(
                f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm'."
            )

    def forward(self, hidden_states: Tensor, emb: Tensor):
        emb = self.linear(self.silu(emb).cast(hidden_states.dtype))