(self, cfg: Config, *, parent: Optional[Module])
| 67 | output_norm: Optional[InstantiableConfig] = None |
| 68 | |
| 69 | def __init__(self, cfg: Config, *, parent: Optional[Module]): |
| 70 | super().__init__(cfg, parent=parent) |
| 71 | cfg = self.config |
| 72 | self._add_child( |
| 73 | "embed_proj", cfg.embed_proj.set(input_dim=cfg.pos_embed_dim, output_dim=cfg.output_dim) |
| 74 | ) |
| 75 | self._add_child( |
| 76 | "output_proj", cfg.output_proj.set(input_dim=cfg.output_dim, output_dim=cfg.output_dim) |
| 77 | ) |
| 78 | if cfg.output_norm is not None: |
| 79 | self._add_child("output_norm", cfg.output_norm.set(input_dim=cfg.output_dim)) |
| 80 | |
| 81 | def dit_sinusoidal_positional_embeddings(self, positions: Tensor): |
| 82 | """DiT Specific Sinusoidal Positional Embeddings. |
nothing calls this directly
no test coverage detected