4cbe889c27f5a0c272c54b4e1065ef3525c2fc09,s2cnn/nn/soft/gpu/so3_fft.py,,_so3_rifft,#Any#Any#Any#Any#,162

Before Change

    device = x.get_device()
    nbatch = x.size(1)

    plan = _setup_fft_plan(b_out, nbatch)
    wigner = _setup_wigner(b_out, nl=b_in, weighted=for_grad, device=device) // [beta, l * m * n] (2 * b_out, nspec)
    cuda_kernel = _setup_so3ifft_cuda_kernel(b_in=b_in, b_out=b_out, nbatch=nbatch, real_output=True)

    output = torch.cuda.FloatTensor(nbatch, 2 * b_out, 2 * b_out, 2 * b_out, 2)
    cuda_kernel(x, wigner, output) // [batch, beta, m, n, complex]

    plan(output, output, 1) // [batch, beta, alpha, gamma, complex]
    output = output[..., 0] // [batch, beta, alpha, gamma]
    return output

After Change

    wigner = _setup_wigner(b_out, nl=b_in, weighted=for_grad, like=x) // [beta, l * m * n] (2 * b_out, nspec)
    cuda_kernel = _setup_so3ifft_cuda_kernel(b_in=b_in, b_out=b_out, nbatch=nbatch, real_output=True)

    output = x.new_empty((nbatch, 2 * b_out, 2 * b_out, 2 * b_out, 2))
    cuda_kernel(x, wigner, output) // [batch, beta, m, n, complex]

    output = torch.ifft(output, 2) * output.size(-2) ** 2  // [batch, beta, alpha, gamma, complex]
    output = output[..., 0] // [batch, beta, alpha, gamma]
    return output

Instances