| 1965 | # Copied from https://github.com/deep-floyd/IF/blob/2f91391f27dd3c468bf174be5805b4cc92980c0b/deepfloyd_if/model/nn.py#L54 |
| 1966 | |
| 1967 | def __init__(self, num_heads, embed_dim, dtype=None): |
| 1968 | super().__init__() |
| 1969 | self.dtype = dtype |
| 1970 | self.positional_embedding = nn.Parameter(torch.randn(1, embed_dim) / embed_dim**0.5) |
| 1971 | self.k_proj = nn.Linear(embed_dim, embed_dim, dtype=self.dtype) |
| 1972 | self.q_proj = nn.Linear(embed_dim, embed_dim, dtype=self.dtype) |
| 1973 | self.v_proj = nn.Linear(embed_dim, embed_dim, dtype=self.dtype) |
| 1974 | self.num_heads = num_heads |
| 1975 | self.dim_per_head = embed_dim // self.num_heads |
| 1976 | |
| 1977 | def forward(self, x): |
| 1978 | bs, length, width = x.size() |