Method forward

pointops2/functions/pointops.py:209–228

input: attn: (M, h), v: (N, h, C//h), index0: (M), index1: (M) output: output: [N, h, C//h]

(ctx, attn, v, index0, index1)

class AttentionStep2(Function):
    @staticmethod
    def forward(ctx, attn, v, index0, index1):
        """
        input: attn: (M, h), v: (N, h, C//h), index0: (M), index1: (M)
        output: output: [N, h, C//h]
        """
        assert attn.is_contiguous() and v.is_contiguous() and index0.is_contiguous() and index1.is_contiguous()

        M, h = attn.shape
        N_q = index0.max().item() + 1
        N_v, h, C_div_h = v.shape
        C = int(C_div_h * h)

        output = torch.cuda.FloatTensor(N_q, h, C//h).zero_()
        pointops_cuda.attention_step2_forward_cuda(N_q, M, h, C, attn, v, index0, index1, output)
        ctx.M = M

        # print("attn[:5,:5]: ", attn[:5, :5])

        ctx.save_for_backward(attn, v, index0, index1)
        return output

    @staticmethod
    def backward(ctx, grad_output):