4cbe889c27f5a0c272c54b4e1065ef3525c2fc09,s2cnn/nn/soft/gpu/so3_fft.py,,_so3_ifft,#Any#Any#Any#Any#,121

Before Change

    :param x: [l * m * n, batch, complex] (b_in (4 b_in**2 - 1) // 3, nbatch, 2)
    :return: [batch, beta, alpha, gamma, complex] (nbatch, 2 b_out, 2 b_out, 2 b_out, 2)
    """
    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)

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=False)

    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]
    return output

Instances