(self, coord)
| 305 | torch.nn.init.trunc_normal_(self.rpe_table, std=0.02) |
| 306 | |
| 307 | def forward(self, coord): |
| 308 | idx = ( |
| 309 | coord.clamp(-self.pos_bnd, self.pos_bnd) # clamp into bnd |
| 310 | + self.pos_bnd # relative position to positive index |
| 311 | + torch.arange(3, device=coord.device) * self.rpe_num # x, y, z stride |
| 312 | ) |
| 313 | out = self.rpe_table.index_select(0, idx.reshape(-1)) |
| 314 | out = out.view(idx.shape + (-1,)).sum(3) |
| 315 | out = out.permute(0, 3, 1, 2) # (N, K, K, H) -> (N, H, K, K) |
| 316 | return out |
| 317 | |
| 318 | |
| 319 | class SerializedAttention(PointModule): |
nothing calls this directly
no outgoing calls
no test coverage detected