Method __init__

axlearn/common/dit.py:159–168

(self, cfg: Config, *, parent: Optional[Module])

        emb: Embedding.Config = Embedding.default_config()

    def __init__(self, cfg: Config, *, parent: Optional[Module]):
        super().__init__(cfg, parent=parent)
        cfg = self.config
        # Ref: https://github.com/facebookresearch/DiT/blob/main/models.py#L74
        # The num_embeddings = cfg.num_classes + (dropout_rate > 0)
        # Additional token is added to label dropout for classifier-free guidance.
        use_drop_rate = int(cfg.dropout_rate > 0)
        self._add_child(
            "emb", cfg.emb.set(dim=cfg.output_dim, num_embeddings=cfg.num_classes + use_drop_rate)
        )

    def _mask_label(self, label: Tensor) -> Tensor:
        """This function is activated only when dropout_rate > 0 and is_training=False.