Method
__init__
(self,
dim,
num_heads,
causal=False,
attn_dropout=0.0,
proj_dropout=0.0)
Source from the content-addressed store, hash-verified
| 53 | class SelfAttention(nn.Module): |
| 54 | |
| 55 | def __init__(self, |
| 56 | dim, |
| 57 | num_heads, |
| 58 | causal=False, |
| 59 | attn_dropout=0.0, |
| 60 | proj_dropout=0.0): |
| 61 | assert dim % num_heads == 0 |
| 62 | super().__init__() |
| 63 | self.dim = dim |
| 64 | self.num_heads = num_heads |
| 65 | self.head_dim = dim // num_heads |
| 66 | self.causal = causal |
| 67 | self.attn_dropout = attn_dropout |
| 68 | self.proj_dropout = proj_dropout |
| 69 | |
| 70 | # layers |
| 71 | self.to_qkv = nn.Linear(dim, dim * 3) |
| 72 | self.proj = nn.Linear(dim, dim) |
| 73 | |
| 74 | def forward(self, x): |
| 75 | """ |
Callers
nothing calls this directly
Tested by
no test coverage detected