Method forward

pointops2/functions/pointops.py:144–164

input: q: (N, h, C//h), k: (N, h, C//h), index0: (M), index1: (M) output: output: [N, h, C//h]

(ctx, q, k, index1, index0_offsets, n_max)

class AttentionStep1_v2(Function):
    @staticmethod
    def forward(ctx, q, k, index1, index0_offsets, n_max):
        """
        input: q: (N, h, C//h), k: (N, h, C//h), index0: (M), index1: (M)
        output: output: [N, h, C//h]
        """
        assert q.is_contiguous() and k.is_contiguous() and index0_offsets.is_contiguous() and index1.is_contiguous()
        assert n_max <= 1024

        N_q, h, C_div_h = q.shape
        N_k = k.shape[0]
        M = index1.shape[0]
        C = int(C_div_h * h)

        output = torch.cuda.FloatTensor(M, h).zero_()
        pointops_cuda.attention_step1_forward_cuda_v2(N_k, M, h, C, n_max, q, k, index0_offsets, index1, output)
        ctx.N_q = N_q
        ctx.N_k = N_k
        ctx.C = C
        ctx.n_max = n_max
        ctx.save_for_backward(q, k, index0_offsets, index1)
        return output

    @staticmethod
    def backward(ctx, grad_output):