How to use the torchkbnufft.KbNufft function in torchkbnufft

dtype = testing_dtype
    norm_tol = testing_tol

    im_size = params_2d['im_size']
    numpoints = params_2d['numpoints']

    x = params_2d['x']
    y = params_2d['y']
    ktraj = params_2d['ktraj']

    for device in device_list:
        x = x.detach().to(dtype=dtype, device=device)
        y = y.detach().to(dtype=dtype, device=device)
        ktraj = ktraj.detach().to(dtype=dtype, device=device)

        kbnufft_ob = KbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)
        adjkbnufft_ob = AdjKbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)

        real_mat, imag_mat = precomp_sparse_mats(ktraj, kbnufft_ob)
        interp_mats = {
            'real_interp_mats': real_mat,
            'imag_interp_mats': imag_mat
        }

        x.requires_grad = True
        y = kbnufft_ob.forward(x, ktraj, interp_mats)

dtype = testing_dtype
    norm_tol = testing_tol

    im_size = params_3d['im_size']
    numpoints = params_3d['numpoints']

    x = params_3d['x']
    y = params_3d['y']
    ktraj = params_3d['ktraj']

    for device in device_list:
        x = x.detach().to(dtype=dtype, device=device)
        y = y.detach().to(dtype=dtype, device=device)
        ktraj = ktraj.detach().to(dtype=dtype, device=device)

        kbnufft_ob = KbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)
        adjkbnufft_ob = AdjKbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)

        real_mat, imag_mat = precomp_sparse_mats(ktraj, kbnufft_ob)
        interp_mats = {
            'real_interp_mats': real_mat,
            'imag_interp_mats': imag_mat
        }

        y.requires_grad = True
        x = adjkbnufft_ob.forward(y, ktraj, interp_mats)

dtype = testing_dtype
    norm_tol = testing_tol

    im_size = params_2d['im_size']
    numpoints = params_2d['numpoints']

    x = params_2d['x']
    y = params_2d['y']
    ktraj = params_2d['ktraj']

    for device in device_list:
        x = x.detach().to(dtype=dtype, device=device)
        y = y.detach().to(dtype=dtype, device=device)
        ktraj = ktraj.detach().to(dtype=dtype, device=device)

        kbnufft_ob = KbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)
        adjkbnufft_ob = AdjKbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)

        y.requires_grad = True
        x = adjkbnufft_ob.forward(y, ktraj)

        ((x ** 2) / 2).sum().backward()
        y_grad = y.grad.clone().detach()

        y_hat = kbnufft_ob.forward(x.clone().detach(), ktraj)

grid_sz = (512, 512)
    n_shift = (128, 128)
    numpoints = 6
    table_oversamp = 2**10
    kbwidth = 2.34
    order = 0
    norm = 'None'

    ob = KbInterpForw(
        im_size=im_sz, grid_size=grid_sz, n_shift=n_shift, numpoints=numpoints,
        table_oversamp=table_oversamp, kbwidth=kbwidth, order=order)
    ob = KbInterpBack(
        im_size=im_sz, grid_size=grid_sz, n_shift=n_shift, numpoints=numpoints,
        table_oversamp=table_oversamp, kbwidth=kbwidth, order=order)

    ob = KbNufft(
        im_size=im_sz, grid_size=grid_sz, n_shift=n_shift, numpoints=numpoints,
        table_oversamp=table_oversamp, kbwidth=kbwidth, order=order, norm=norm)
    ob = AdjKbNufft(
        im_size=im_sz, grid_size=grid_sz, n_shift=n_shift, numpoints=numpoints,
        table_oversamp=table_oversamp, kbwidth=kbwidth, order=order, norm=norm)

    ob = MriSenseNufft(
        smap=smap, im_size=im_sz, grid_size=grid_sz, n_shift=n_shift, numpoints=numpoints,
        table_oversamp=table_oversamp, kbwidth=kbwidth, order=order, norm=norm)
    ob = AdjMriSenseNufft(
        smap=smap, im_size=im_sz, grid_size=grid_sz, n_shift=n_shift, numpoints=numpoints,
        table_oversamp=table_oversamp, kbwidth=kbwidth, order=order, norm=norm)

    # test 2d tuple inputs
    im_sz = (256, 256)
    smap = torch.randn(*((1,) + im_sz))

dtype = testing_dtype
    norm_tol = testing_tol

    im_size = params_3d['im_size']
    numpoints = params_3d['numpoints']

    x = params_3d['x']
    y = params_3d['y']
    ktraj = params_3d['ktraj']

    for device in device_list:
        x = x.detach().to(dtype=dtype, device=device)
        y = y.detach().to(dtype=dtype, device=device)
        ktraj = ktraj.detach().to(dtype=dtype, device=device)

        kbnufft_ob = KbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)
        adjkbnufft_ob = AdjKbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)

        x.requires_grad = True
        y = kbnufft_ob.forward(x, ktraj)

        ((y ** 2) / 2).sum().backward()
        x_grad = x.grad.clone().detach()

        x_hat = adjkbnufft_ob.forward(y.clone().detach(), ktraj)

dtype = testing_dtype
    norm_tol = testing_tol

    im_size = params_2d['im_size']
    numpoints = params_2d['numpoints']

    x = params_2d['x']
    y = params_2d['y']
    ktraj = params_2d['ktraj']

    for device in device_list:
        x = x.detach().to(dtype=dtype, device=device)
        y = y.detach().to(dtype=dtype, device=device)
        ktraj = ktraj.detach().to(dtype=dtype, device=device)

        kbnufft_ob = KbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)
        adjkbnufft_ob = AdjKbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)

        x_forw = kbnufft_ob(x, ktraj)
        y_back = adjkbnufft_ob(y, ktraj)

        inprod1 = inner_product(y, x_forw, dim=2)
        inprod2 = inner_product(y_back, x, dim=2)

        assert torch.norm(inprod1 - inprod2) < norm_tol

dtype = testing_dtype
    norm_tol = testing_tol

    im_size = params_3d['im_size']
    numpoints = params_3d['numpoints']

    x = params_3d['x']
    y = params_3d['y']
    ktraj = params_3d['ktraj']

    for device in device_list:
        x = x.detach().to(dtype=dtype, device=device)
        y = y.detach().to(dtype=dtype, device=device)
        ktraj = ktraj.detach().to(dtype=dtype, device=device)

        kbnufft_ob = KbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)
        adjkbnufft_ob = AdjKbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)

        real_mat, imag_mat = precomp_sparse_mats(ktraj, kbnufft_ob)
        interp_mats = {
            'real_interp_mats': real_mat,
            'imag_interp_mats': imag_mat
        }

        x_forw = kbnufft_ob(x, ktraj, interp_mats)
        y_back = adjkbnufft_ob(y, ktraj, interp_mats)

dtype = testing_dtype
    norm_tol = testing_tol

    im_size = params_2d['im_size']
    numpoints = params_2d['numpoints']

    x = params_2d['x']
    y = params_2d['y']
    ktraj = params_2d['ktraj']

    for device in device_list:
        x = x.detach().to(dtype=dtype, device=device)
        y = y.detach().to(dtype=dtype, device=device)
        ktraj = ktraj.detach().to(dtype=dtype, device=device)

        kbnufft_ob = KbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)
        adjkbnufft_ob = AdjKbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)

        real_mat, imag_mat = precomp_sparse_mats(ktraj, kbnufft_ob)
        interp_mats = {
            'real_interp_mats': real_mat,
            'imag_interp_mats': imag_mat
        }

        x_forw = kbnufft_ob(x, ktraj, interp_mats)
        y_back = adjkbnufft_ob(y, ktraj, interp_mats)

dtype = testing_dtype
    norm_tol = testing_tol

    im_size = params_3d['im_size']
    numpoints = params_3d['numpoints']

    x = params_3d['x']
    y = params_3d['y']
    ktraj = params_3d['ktraj']

    for device in device_list:
        x = x.detach().to(dtype=dtype, device=device)
        y = y.detach().to(dtype=dtype, device=device)
        ktraj = ktraj.detach().to(dtype=dtype, device=device)

        kbnufft_ob = KbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)
        adjkbnufft_ob = AdjKbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)

        real_mat, imag_mat = precomp_sparse_mats(ktraj, kbnufft_ob)
        interp_mats = {
            'real_interp_mats': real_mat,
            'imag_interp_mats': imag_mat
        }

        x.requires_grad = True
        y = kbnufft_ob.forward(x, ktraj, interp_mats)

dtype = testing_dtype
    norm_tol = testing_tol

    im_size = params_2d['im_size']
    numpoints = params_2d['numpoints']

    x = params_2d['x']
    y = params_2d['y']
    ktraj = params_2d['ktraj']

    for device in device_list:
        x = x.detach().to(dtype=dtype, device=device)
        y = y.detach().to(dtype=dtype, device=device)
        ktraj = ktraj.detach().to(dtype=dtype, device=device)

        kbnufft_ob = KbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)
        adjkbnufft_ob = AdjKbNufft(
            im_size=im_size,
            numpoints=numpoints
        ).to(dtype=dtype, device=device)

        x.requires_grad = True
        y = kbnufft_ob.forward(x, ktraj)

        ((y ** 2) / 2).sum().backward()
        x_grad = x.grad.clone().detach()

        x_hat = adjkbnufft_ob.forward(y.clone().detach(), ktraj)