input: q: (N, h, C//h), k: (N, h, C//h), index0: (M), index1: (M) output: output: [N, h, C//h]
(ctx, q, k, index0, index1)
| 82 | class AttentionStep1(Function): |
| 83 | @staticmethod |
| 84 | def forward(ctx, q, k, index0, index1): |
| 85 | """ |
| 86 | input: q: (N, h, C//h), k: (N, h, C//h), index0: (M), index1: (M) |
| 87 | output: output: [N, h, C//h] |
| 88 | """ |
| 89 | assert q.is_contiguous() and k.is_contiguous() and index0.is_contiguous() and index1.is_contiguous() |
| 90 | |
| 91 | N_q, h, C_div_h = q.shape |
| 92 | N_k = k.shape[0] |
| 93 | M = index0.shape[0] |
| 94 | C = int(C_div_h * h) |
| 95 | |
| 96 | output = torch.cuda.FloatTensor(M, h).zero_() |
| 97 | pointops_cuda.attention_step1_forward_cuda(N_k, M, h, C, q, k, index0, index1, output) |
| 98 | ctx.N_q = N_q |
| 99 | ctx.N_k = N_k |
| 100 | ctx.C = C |
| 101 | ctx.save_for_backward(q, k, index0, index1) |
| 102 | return output |
| 103 | |
| 104 | @staticmethod |
| 105 | def backward(ctx, grad_output): |
nothing calls this directly
no outgoing calls
no test coverage detected