jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_geometry(
            self,
            fH,
            fW,
            lidar2img,
            img_metas,
    ):
        B, N, _, _ = lidar2img.shape
        device = lidar2img.device
        if self.frustum == None:
            self.frustum = self.create_frustum(fH, fW, img_metas)
            self.frustum = self.frustum.to(device)
            # self.D = self.frustum.shape[0]

        points = self.frustum.view(1, 1, self.D, fH, fW, 3) \
            .repeat(B, N, 1, 1, 1, 1)
        lidar2img = lidar2img.view(B, N, 1, 1, 1, 4, 4)
        # img2lidar = torch.inverse(lidar2img)
        points = torch.cat(
            (points, torch.ones_like(points[..., :1])), -1)
        points = torch.linalg.solve(lidar2img.to(torch.float32),
                                    points.unsqueeze(-1).to(torch.float32)).squeeze(-1)
        # points = torch.matmul(img2lidar.to(torch.float32),
        #                       points.unsqueeze(-1).to(torch.float32)).squeeze(-1)
        eps = 1e-5
        points = points[..., 0:3] / torch.maximum(
            points[..., 3:4], torch.ones_like(points[..., 3:4]) * eps)

        return points
//  [python]
def get_geometry(
            self,
            fH,
            fW,
            lidar2img,
            img_metas,
    ):
        B, N, _, _ = lidar2img.shape
        device = lidar2img.device
        # import pdb;pdb.set_trace()
        if self.frustum == None:
            self.frustum = self.create_frustum(fH, fW, img_metas)
            self.frustum = self.frustum.to(device)
            # self.D = self.frustum.shape[0]

        points = self.frustum.view(1, 1, self.D, fH, fW, 3) \
            .repeat(B, N, 1, 1, 1, 1)
        lidar2img = lidar2img.view(B, N, 1, 1, 1, 4, 4)
        # img2lidar = torch.inverse(lidar2img)
        points = torch.cat(
            (points, torch.ones_like(points[..., :1])), -1)
        points = torch.solve(points.unsqueeze(-1).to(torch.float32),
                                    lidar2img.to(torch.float32))[0].squeeze(-1)
        # points = torch.matmul(img2lidar.to(torch.float32),
        #                       points.unsqueeze(-1).to(torch.float32)).squeeze(-1)
        # import pdb;pdb.set_trace()
        eps = 1e-5
        points = points[..., 0:3] / torch.maximum(
            points[..., 3:4], torch.ones_like(points[..., 3:4]) * eps)

        return points
//  [python]
def get_geometry(
            self,
            fH,
            fW,
            lidar2img,
            img_metas,
    ):
        B, N, _, _ = lidar2img.shape
        device = lidar2img.device
        # import pdb;pdb.set_trace()
        if self.frustum == None:
            self.frustum = self.create_frustum(fH, fW, img_metas)
            self.frustum = self.frustum.to(device)
            # self.D = self.frustum.shape[0]

        points = self.frustum.view(1, 1, self.D, fH, fW, 3) \
            .repeat(B, N, 1, 1, 1, 1)
        lidar2img = lidar2img.view(B, N, 1, 1, 1, 4, 4)
        # img2lidar = torch.inverse(lidar2img)
        points = torch.cat(
            (points, torch.ones_like(points[..., :1])), -1)
        points = torch.linalg.solve(lidar2img.to(torch.float32),
                                    points.unsqueeze(-1).to(torch.float32)).squeeze(-1)
        # points = torch.matmul(img2lidar.to(torch.float32),
        #                       points.unsqueeze(-1).to(torch.float32)).squeeze(-1)
        # import pdb;pdb.set_trace()
        eps = 1e-5
        points = points[..., 0:3] / torch.maximum(
            points[..., 3:4], torch.ones_like(points[..., 3:4]) * eps)

        return points
//  [python]
def _crop_y(y, sett):
    """ Crop output images FOV to a fixed dimension

    Args:
        y (_output()): _output data.

    Returns:
        y (_output()): Cropped output data.

    """
    if not sett.crop:
        return y
    device = sett.device
    # Output image information
    mat_y = y[0].mat
    vx_y = voxel_size(mat_y)
    # Define cropped FOV
    mat_mu, dim_mu = _bb_atlas('atlas_t1',
        fov=sett.fov, dtype=torch.float64, device=device)
    # Modulate atlas with voxel size
    mat_vx = torch.diag(torch.cat((
        vx_y, torch.ones(1, dtype=torch.float64, device=device))))
    mat_mu = mat_mu.mm(mat_vx)
    dim_mu = mat_vx[:3, :3].inverse().mm(dim_mu[:, None]).floor().squeeze()
    # Make output grid
    M = torch.solve(mat_mu, mat_y)[0].type(y[0].dat.dtype)
    grid = affine_grid(M, dim_mu)[None, ...]
    # Crop
    for c in range(len(y)):
        y[c].dat = grid_pull(y[c].dat[None, None, ...], grid,
                             bound='zero', extrapolate=False,
                             interpolation=0)[0, 0, ...]
        # Do labels?
        if y[c].label is not None:
            y[c].label = grid_pull(y[c].label[None, None, ...], grid,
                                   bound='zero', extrapolate=False,
                                   interpolation=0)[0, 0, ...]
        y[c].mat = mat_mu
        y[c].dim = tuple(dim_mu.int().tolist())

    return y
//  [python]
def _crop_y(y, sett):
    """ Crop output images FOV to a fixed dimension

    Args:
        y (_output()): _output data.

    Returns:
        y (_output()): Cropped output data.

    """
    if not sett.crop:
        return y
    device = sett.device
    # Output image information
    mat_y = y[0].mat
    vx_y = voxel_size(mat_y)
    # Define cropped FOV
    mat_mu, dim_mu = _bb_atlas('atlas_t1',
        fov=sett.fov, dtype=torch.float64, device=device)
    # Modulate atlas with voxel size
    mat_vx = torch.diag(torch.cat((
        vx_y, torch.ones(1, dtype=torch.float64, device=device))))
    mat_mu = mat_mu.mm(mat_vx)
    dim_mu = mat_vx[:3, :3].inverse().mm(dim_mu[:, None]).floor().squeeze()
    # Make output grid
    M = torch.linalg.solve(mat_y, mat_mu).type(y[0].dat.dtype)
    grid = affine_grid(M, dim_mu)[None, ...]
    # Crop
    for c in range(len(y)):
        y[c].dat = grid_pull(y[c].dat[None, None, ...], grid,
                             bound='zero', extrapolate=False,
                             interpolation=0)[0, 0, ...]
        # Do labels?
        if y[c].label is not None:
            y[c].label = grid_pull(y[c].label[None, None, ...], grid,
                                   bound='zero', extrapolate=False,
                                   interpolation=0)[0, 0, ...]
        y[c].mat = mat_mu
        y[c].dim = tuple(dim_mu.int().tolist())

    return y
//  [python]
def update_precond_UVd_math_(U, V, d, v, h, step, tiny):
    # type: (Tensor, Tensor, Tensor, Tensor, Tensor, float, float) -> None
    """
    Update preconditioner Q = (I + U*V')*diag(d) with (vector, Hessian-vector product) = (v, h).
    State variables U, V and d are updated inplace. 
                               
    U, V, d, v, and h are either matrices or column vectors.  
    """
    # balance the numerical dynamic ranges of U and V; optional 
    if torch.rand([]) < 0.01:
        normU = torch.linalg.vector_norm(U)
        normV = torch.linalg.vector_norm(V)
        rho = torch.sqrt(normU/normV)
        U.div_(rho)
        V.mul_(rho)

    Qh = IpUVtmatvec(U, V, d*h)
    Ph = d*IpUVtmatvec(V, U, Qh)
    
    # invQtv = IpUVtsolve(V, U, v/d)
    # invPv = IpUVtsolve(U, V, invQtv)/d
    VtU = V.t().mm(U)
    I = torch.eye(VtU.size(dim=0), dtype=VtU.dtype, device=VtU.device)
    IpVtU = I + VtU
    invQtv = v/d
    # torch's linalg.solve is slow for small matrix
    invQtv = invQtv - V.mm(torch.solve(U.t().mm(invQtv), IpVtU.t())[0])  
    invPv  = invQtv - U.mm(torch.solve(V.t().mm(invQtv), IpVtU)[0])
    invPv = invPv/d

    nablaD = Ph*h - v*invPv
    mu = step/(torch.max(torch.abs(nablaD)) + tiny)
    #d = d - mu*d*nablaD
    d.sub_(mu*d*nablaD)
    
    # update either U or V, not both at the same time
    a, b = Qh, invQtv
    if torch.rand([]) < 0.5:
        # nablaU = Qh.mm(Qh.t().mm(V)) - invQtv.mm(invQtv.t().mm(V))
        # mu = step/(norm_UVt(nablaU, V) + _tiny)
        # U = U - mu*(nablaU + nablaU.mm(V.t().mm(U)))
        atV = a.t().mm(V)
        atVVt = atV.mm(V.t())
        btV = b.t().mm(V)
        btVVt = btV.mm(V.t())
        norm = torch.sqrt(torch.abs( (a.t().mm(a))*(atVVt.mm(atVVt.t())) # abs to avoid sqrt(-0.0) 
                                    +(b.t().mm(b))*(btVVt.mm(btVVt.t())) 
                                  -2*(a.t().mm(b))*(atVVt.mm(btVVt.t())) ))
        mu = step/(norm + tiny)
        # U = U - mu*( a.mm(atV.mm(IpVtU)) 
        #             -b.mm(btV.mm(IpVtU)) )
        U.sub_(mu*( a.mm(atV.mm(IpVtU)) 
                   -b.mm(btV.mm(IpVtU)) ))
    else:
        # nablaV = Qh.mm(Qh.t().mm(U)) - invQtv.mm(invQtv.t().mm(U))
        # mu = step/(norm_UVt(U, nablaV) + _tiny)
        # V = V - mu*(nablaV + V.mm(U.t().mm(nablaV)))
        atU = a.t().mm(U)
        btU = b.t().mm(U)
        UUta = U.mm(atU.t())
        UUtb = U.mm(btU.t())
        norm = torch.sqrt(torch.abs( (UUta.t().mm(UUta))*(a.t().mm(a)) # abs to avoid sqrt(-0.0)
                                    +(UUtb.t().mm(UUtb))*(b.t().mm(b))
                                  -2*(UUta.t().mm(UUtb))*(a.t().mm(b)) ))
        mu = step/(norm + tiny)
        # V = V - mu*( (a + V.mm(atU.t())).mm(atU) 
        #             -(b + V.mm(btU.t())).mm(btU) )
        V.sub_(mu*( (a + V.mm(atU.t())).mm(atU) 
                   -(b + V.mm(btU.t())).mm(btU) ))
//  [python]
def update_precond_UVd_math_(U, V, d, v, h, step, tiny):
    # type: (Tensor, Tensor, Tensor, Tensor, Tensor, float, float) -> None
    """
    Update preconditioner Q = (I + U*V')*diag(d) with (vector, Hessian-vector product) = (v, h).
    State variables U, V and d are updated inplace. 
                               
    U, V, d, v, and h are either matrices or column vectors.  
    """
    # balance the numerical dynamic ranges of U and V; optional 
    if torch.rand([]) < 0.01:
        normU = torch.linalg.vector_norm(U)
        normV = torch.linalg.vector_norm(V)
        rho = torch.sqrt(normU/normV)
        U.div_(rho)
        V.mul_(rho)

    Qh = IpUVtmatvec(U, V, d*h)
    Ph = d*IpUVtmatvec(V, U, Qh)
    
    # invQtv = IpUVtsolve(V, U, v/d)
    # invPv = IpUVtsolve(U, V, invQtv)/d
    VtU = V.t().mm(U)
    I = torch.eye(VtU.size(dim=0), dtype=VtU.dtype, device=VtU.device)
    IpVtU = I + VtU
    invQtv = v/d
    # torch's linalg.solve is slow for small matrix
    invQtv = invQtv - V.mm(torch.linalg.solve(IpVtU.t(), U.t().mm(invQtv)))  
    invPv  = invQtv - U.mm(torch.linalg.solve(IpVtU,     V.t().mm(invQtv)))
    invPv = invPv/d

    nablaD = Ph*h - v*invPv
    mu = step/(torch.max(torch.abs(nablaD)) + tiny)
    #d = d - mu*d*nablaD
    d.sub_(mu*d*nablaD)
    
    # update either U or V, not both at the same time
    a, b = Qh, invQtv
    if torch.rand([]) < 0.5:
        # nablaU = Qh.mm(Qh.t().mm(V)) - invQtv.mm(invQtv.t().mm(V))
        # mu = step/(norm_UVt(nablaU, V) + _tiny)
        # U = U - mu*(nablaU + nablaU.mm(V.t().mm(U)))
        atV = a.t().mm(V)
        atVVt = atV.mm(V.t())
        btV = b.t().mm(V)
        btVVt = btV.mm(V.t())
        norm = torch.sqrt(torch.abs( (a.t().mm(a))*(atVVt.mm(atVVt.t())) # abs to avoid sqrt(-0.0) 
                                    +(b.t().mm(b))*(btVVt.mm(btVVt.t())) 
                                  -2*(a.t().mm(b))*(atVVt.mm(btVVt.t())) ))
        mu = step/(norm + tiny)
        # U = U - mu*( a.mm(atV.mm(IpVtU)) 
        #             -b.mm(btV.mm(IpVtU)) )
        U.sub_(mu*( a.mm(atV.mm(IpVtU)) 
                   -b.mm(btV.mm(IpVtU)) ))
    else:
        # nablaV = Qh.mm(Qh.t().mm(U)) - invQtv.mm(invQtv.t().mm(U))
        # mu = step/(norm_UVt(U, nablaV) + _tiny)
        # V = V - mu*(nablaV + V.mm(U.t().mm(nablaV)))
        atU = a.t().mm(U)
        btU = b.t().mm(U)
        UUta = U.mm(atU.t())
        UUtb = U.mm(btU.t())
        norm = torch.sqrt(torch.abs( (UUta.t().mm(UUta))*(a.t().mm(a)) # abs to avoid sqrt(-0.0)
                                    +(UUtb.t().mm(UUtb))*(b.t().mm(b))
                                  -2*(UUta.t().mm(UUtb))*(a.t().mm(b)) ))
        mu = step/(norm + tiny)
        # V = V - mu*( (a + V.mm(atU.t())).mm(atU) 
        #             -(b + V.mm(btU.t())).mm(btU) )
        V.sub_(mu*( (a + V.mm(atU.t())).mm(atU) 
                   -(b + V.mm(btU.t())).mm(btU) ))
//  [python]
def _init_y_label(x, y, sett):
    """Make initial guess of labels.
    """
    dim_y = y[0].dim
    mat_y = y[0].mat
    for c in range(len(x)):
        n = 0
        if x[c][n].label is not None:
            # Make output grid
            mat = torch.solve(mat_y, x[c][n].mat)[0]  # mat_x\mat_y
            grid = affine_grid(mat.type(x[c][n].dat.dtype), dim_y)
            # Do resampling
            y[c].label = _warp_label(x[c][n].label[0], grid)

    return y
//  [python]
def _init_y_label(x, y, sett):
    """Make initial guess of labels.
    """
    dim_y = y[0].dim
    mat_y = y[0].mat
    for c in range(len(x)):
        n = 0
        if x[c][n].label is not None:
            # Make output grid
            mat = torch.linalg.solve(x[c][n].mat, mat_y)  # mat_x\mat_y
            grid = affine_grid(mat.type(x[c][n].dat.dtype), dim_y)
            # Do resampling
            y[c].label = _warp_label(x[c][n].label[0], grid)

    return y
//  [python]
def _init_y_dat(x, y, sett):
    """ Make initial guesses of reconstucted image(s) using b-spline interpolation,
        with averaging if more than one observation per channel.
    """
    dim_y = y[0].dim
    mat_y = y[0].mat
    for c in range(len(x)):
        dat_y = torch.zeros(dim_y, dtype=torch.float32, device=sett.device)
        num_x = len(x[c])
        sm    = torch.zeros_like(dat_y)
        for n in range(num_x):
            # Get image data
            dat = x[c][n].dat[None, None, ...]
            # Make output grid
            mat = torch.solve(mat_y, x[c][n].mat)[0]  # mat_x\mat_y            
            grid = affine_grid(mat.type(dat.dtype), dim_y)
            # Do resampling
            mn = torch.min(dat)
            mx = torch.max(dat)
            dat = grid_pull(dat, grid[None, ...],
                bound='zero', extrapolate=False, interpolation=1)
            dat[dat < mn] = mn
            dat[dat > mx] = mx
            sm = sm + (dat[0, 0, ...] > 0)
            dat_y = dat_y + dat[0, 0, ...]
        sm[sm == 0] = 1.0
        y[c].dat = dat_y / sm

    return y
//  [python]
def _init_y_dat(x, y, sett):
    """ Make initial guesses of reconstucted image(s) using b-spline interpolation,
        with averaging if more than one observation per channel.
    """
    dim_y = y[0].dim
    mat_y = y[0].mat
    for c in range(len(x)):
        dat_y = torch.zeros(dim_y, dtype=torch.float32, device=sett.device)
        num_x = len(x[c])
        sm    = torch.zeros_like(dat_y)
        for n in range(num_x):
            # Get image data
            dat = x[c][n].dat[None, None, ...]
            # Make output grid
            mat = torch.linalg.solve(x[c][n].mat, mat_y)  # mat_x\mat_y            
            grid = affine_grid(mat.type(dat.dtype), dim_y)
            # Do resampling
            mn = torch.min(dat)
            mx = torch.max(dat)
            dat = grid_pull(dat, grid[None, ...],
                bound='zero', extrapolate=False, interpolation=1)
            dat[dat < mn] = mn
            dat[dat > mx] = mx
            sm = sm + (dat[0, 0, ...] > 0)
            dat_y = dat_y + dat[0, 0, ...]
        sm[sm == 0] = 1.0
        y[c].dat = dat_y / sm

    return y
//  [python]
def moving_least_square(x: torch.Tensor, y: torch.Tensor, w: torch.Tensor):
    # 1-D MLS: x: (..., N), y: (..., N), w: (..., N)
    p = torch.stack([torch.ones_like(x), x], dim=-2)             # (..., 2, N)
    M = p @ (w[..., :, None] * p.transpose(-2, -1))
    a = torch.solve(M, (p @ (w * y)[..., :, None]))[0]
    a = a.squeeze(-1)
    return a
//  [python]
def moving_least_square(x: torch.Tensor, y: torch.Tensor, w: torch.Tensor):
    # 1-D MLS: x: (..., N), y: (..., N), w: (..., N)
    p = torch.stack([torch.ones_like(x), x], dim=-2)             # (..., 2, N)
    M = p @ (w[..., :, None] * p.transpose(-2, -1))
    a = torch.linalg.solve(M, (p @ (w * y)[..., :, None]))
    a = a.squeeze(-1)
    return a
//  [python]
def compute_input_latents(decoder, r, t, ctf, input_ims):
    bs = r.shape[0]
    E = torch.stack(
        bs*[generate_latent_units(decoder.latent_dim)]).to(decoder.device)
    ims, _, _ = decoder(E, r, t)
    input_ims = apply_ctf(input_ims, ctf)
    y = input_ims.flatten(start_dim=1)
    A = ims.flatten(start_dim=2)
    ATA = A.transpose(1, 2)
    lat = torch.solve(y, ATA)[0]
    return lat
//  [python]
def compute_input_latents(decoder, r, t, ctf, input_ims):
    bs = r.shape[0]
    E = torch.stack(
        bs*[generate_latent_units(decoder.latent_dim)]).to(decoder.device)
    ims, _, _ = decoder(E, r, t)
    input_ims = apply_ctf(input_ims, ctf)
    y = input_ims.flatten(start_dim=1)
    A = ims.flatten(start_dim=2)
    ATA = A.transpose(1, 2)
    lat = torch.linalg.solve(ATA, y)
    return lat
//  [python]
use_cuda = torch.cuda.is_available()

if torch.__version__ >= "1.8":
    torchsolve = lambda A, B: torch.solve(B, A)[0]
else:
    torchsolve = lambda A, B: torch.solve(B, A)[0]
//  [python]
use_cuda = torch.cuda.is_available()

if torch.__version__ >= "1.8":
    torchsolve = lambda A, B: torch.linalg.solve(A, B)
else:
    torchsolve = lambda A, B: torch.solve(B, A)[0]
//  [python]
def _init_reg(x, sett):
    """ Initialise registration.

    """
    # Total number of observations
    N = sum([len(xn) for xn in x])
    # Set rigid affine basis
    sett.rigid_basis = affine_basis(
        group='SE', device=sett.device, dtype=torch.float64)
    fix =sett.fix  # Fixed image index

    # Make input for nitorch affine align
    imgs = []
    for c in range(len(x)):
        for n in range(len(x[c])):
            imgs.append([x[c][n].dat, x[c][n].mat])

    if sett.do_coreg and N > 1:
        # Align images, pairwise, to fixed image (fix)
        t0 = _print_info('init-reg', sett, 'co', 'begin', N)
        mat_a = affine_align(imgs, **sett.coreg_params, fix=fix, device=sett.device)[1]
        sett.mat_coreg = mat_a
        # Apply coreg transform
        i = 0
        for c in range(len(x)):
            for n in range(len(x[c])):
                imgs[i][1] = torch.solve(mat_a[i, ...], imgs[i][1])[0]
                i += 1        
        _print_info('init-reg', sett, 'co', 'finished', N, t0)

    if sett.do_atlas_align:
        # Align fixed image to atlas space, and apply transformation to
        # all images
        t0 = _print_info('init-reg', sett, 'atlas', 'begin', N)
        imgs1 = [imgs[fix]]
        _, mat_a, _, mat_cso = atlas_align(imgs1, rigid=sett.atlas_rigid, device=sett.device)
        sett.mat_atlas = mat_a
        _print_info('init-reg', sett, 'atlas', 'finished', N, t0)
        # Apply atlas registration transform
        i = 0
        for c in range(len(x)):
            for n in range(len(x[c])):
                imgs[i][1] = torch.solve(mat_a, imgs[i][1])[0]
                i += 1

    # Modify image affine (label uses the same as the image, so no need to modify that one)
    i = 0
    for c in range(len(x)):
        for n in range(len(x[c])):
            x[c][n].mat = imgs[i][1]
            i += 1

    # Init rigid parameters (for unified rigid registration)
    for c in range(len(x)):  # Loop over channels
        for n in range(len(x[c])):  # Loop over observations of channel c
            x[c][n].rigid_q = torch.zeros(sett.rigid_basis.shape[0],
                device=sett.device, dtype=torch.float64)

    return x, sett
//  [python]
def _init_reg(x, sett):
    """ Initialise registration.

    """
    # Total number of observations
    N = sum([len(xn) for xn in x])
    # Set rigid affine basis
    sett.rigid_basis = affine_basis(
        group='SE', device=sett.device, dtype=torch.float64)
    fix =sett.fix  # Fixed image index

    # Make input for nitorch affine align
    imgs = []
    for c in range(len(x)):
        for n in range(len(x[c])):
            imgs.append([x[c][n].dat, x[c][n].mat])

    if sett.do_coreg and N > 1:
        # Align images, pairwise, to fixed image (fix)
        t0 = _print_info('init-reg', sett, 'co', 'begin', N)
        mat_a = affine_align(imgs, **sett.coreg_params, fix=fix, device=sett.device)[1]
        sett.mat_coreg = mat_a
        # Apply coreg transform
        i = 0
        for c in range(len(x)):
            for n in range(len(x[c])):
                imgs[i][1] = torch.linalg.solve(mat_a[i, ...], imgs[i][1])
                i += 1        
        _print_info('init-reg', sett, 'co', 'finished', N, t0)

    if sett.do_atlas_align:
        # Align fixed image to atlas space, and apply transformation to
        # all images
        t0 = _print_info('init-reg', sett, 'atlas', 'begin', N)
        imgs1 = [imgs[fix]]
        _, mat_a, _, mat_cso = atlas_align(imgs1, rigid=sett.atlas_rigid, device=sett.device)
        sett.mat_atlas = mat_a
        _print_info('init-reg', sett, 'atlas', 'finished', N, t0)
        # Apply atlas registration transform
        i = 0
        for c in range(len(x)):
            for n in range(len(x[c])):
                imgs[i][1] = torch.linalg.solve(mat_a, imgs[i][1])
                i += 1

    # Modify image affine (label uses the same as the image, so no need to modify that one)
    i = 0
    for c in range(len(x)):
        for n in range(len(x[c])):
            x[c][n].mat = imgs[i][1]
            i += 1

    # Init rigid parameters (for unified rigid registration)
    for c in range(len(x)):  # Loop over channels
        for n in range(len(x[c])):  # Loop over observations of channel c
            x[c][n].rigid_q = torch.zeros(sett.rigid_basis.shape[0],
                device=sett.device, dtype=torch.float64)

    return x, sett
//  [python]
def get_cls_free_guided_flow(self, xt, alpha, logits, logits_cond,):
        B, L, K = xt.shape
        probs = torch.nn.functional.softmax(logits, dim=-1)
        probs_cond = torch.nn.functional.softmax(logits_cond, dim=-1)


        cond_scores_mats = ((alpha - 1) * (torch.eye(self.model.alphabet_size).to(xt)[None, :] / xt[..., None]))  # [B, L, K, K]

        cond_scores_mats = cond_scores_mats - cond_scores_mats.mean(2)[:, :, None, :]  # [B, L, K, K] now the columns sum up to 0

        score = torch.einsum('ijkl,ijl->ijk', cond_scores_mats, probs)  # [B, L, K] add up the columns of conditional flow scores weighted by the predicted probability of each corner
        score_cond = torch.einsum('ijkl,ijl->ijk', cond_scores_mats, probs_cond)  # [B, L, K] add up the columns of conditional flow scores weighted by the predicted probability of each corner
        score_guided = (1 - self.args.guidance_scale) * score + self.args.guidance_scale * score_cond

        Q_mats = cond_scores_mats.clone()  # [B, L, K, K]
        Q_mats[:, :, -1, :] = torch.ones((B, L, K))  # [B, L, K, K]
        score_guided_ = score_guided.clone()  # [B, L, K]
        score_guided_[:, :, -1] = torch.ones(B, L)  # [B, L, K]
        flow_guided = torch.solve(score_guided_, Q_mats)[0]  # [B, L, K]
        return flow_guided
//  [python]
def get_cls_free_guided_flow(self, xt, alpha, logits, logits_cond,):
        B, L, K = xt.shape
        probs = torch.nn.functional.softmax(logits, dim=-1)
        probs_cond = torch.nn.functional.softmax(logits_cond, dim=-1)


        cond_scores_mats = ((alpha - 1) * (torch.eye(self.model.alphabet_size).to(xt)[None, :] / xt[..., None]))  # [B, L, K, K]

        cond_scores_mats = cond_scores_mats - cond_scores_mats.mean(2)[:, :, None, :]  # [B, L, K, K] now the columns sum up to 0

        score = torch.einsum('ijkl,ijl->ijk', cond_scores_mats, probs)  # [B, L, K] add up the columns of conditional flow scores weighted by the predicted probability of each corner
        score_cond = torch.einsum('ijkl,ijl->ijk', cond_scores_mats, probs_cond)  # [B, L, K] add up the columns of conditional flow scores weighted by the predicted probability of each corner
        score_guided = (1 - self.args.guidance_scale) * score + self.args.guidance_scale * score_cond

        Q_mats = cond_scores_mats.clone()  # [B, L, K, K]
        Q_mats[:, :, -1, :] = torch.ones((B, L, K))  # [B, L, K, K]
        score_guided_ = score_guided.clone()  # [B, L, K]
        score_guided_[:, :, -1] = torch.ones(B, L)  # [B, L, K]
        flow_guided = torch.linalg.solve(Q_mats, score_guided_)  # [B, L, K]
        return flow_guided
//  [python]
def forward(self, q: Tensor, v: Tensor, u: Tensor) -> Tuple[Tensor, Tensor]:
        """Evaluates Lagrangian dynamics terms at given state and input.

        Args:
            q: (\*, n_q) configuration batch.
            v: (\*, n_v) velocity batch.
            u: (\*, n_u) input batch.

        Returns:
            (\*, n_v, n_v) mass matrix batch M(q)
            (\*, n_v) Lagrangian contact-free acceleration inv(M(q)) F(q)
        """
        # Pylint bug: cannot recognize instance attributes as Callable.
        # pylint: disable=not-callable
        assert self.mass_matrix is not None
        assert self.lagrangian_forces is not None
        inertia = InertialParameterConverter.pi_cm_to_drake_spatial_inertia(
            self.pi_cm())
        inertia = inertia.expand(q.shape[:-1] + inertia.shape)

        M = self.mass_matrix(q, inertia)
        non_contact_acceleration = torch.solve(
            self.lagrangian_forces(q, v, u, inertia), M)[0]
        return M, non_contact_acceleration
//  [python]
def forward(self, q: Tensor, v: Tensor, u: Tensor) -> Tuple[Tensor, Tensor]:
        """Evaluates Lagrangian dynamics terms at given state and input.

        Args:
            q: (\*, n_q) configuration batch.
            v: (\*, n_v) velocity batch.
            u: (\*, n_u) input batch.

        Returns:
            (\*, n_v, n_v) mass matrix batch M(q)
            (\*, n_v) Lagrangian contact-free acceleration inv(M(q)) F(q)
        """
        # Pylint bug: cannot recognize instance attributes as Callable.
        # pylint: disable=not-callable
        assert self.mass_matrix is not None
        assert self.lagrangian_forces is not None
        inertia = InertialParameterConverter.pi_cm_to_drake_spatial_inertia(
            self.pi_cm())
        inertia = inertia.expand(q.shape[:-1] + inertia.shape)

        M = self.mass_matrix(q, inertia)
        non_contact_acceleration = torch.linalg.solve(
            M, self.lagrangian_forces(q, v, u, inertia))
        return M, non_contact_acceleration
//  [python]
def test_complex_valued_inverse(self):
        a_dim = 3
        a_mat = torch.tensor([[1, 2, 3], [2, 1, 2], [3, 2, 1]], dtype=torch.complex64)

        def a_op(x):
            return a_mat @ x

        cg_solver = ConjugateGradient(a_op, num_iter=100)
        y = torch.tensor([1, 2, 3], dtype=torch.complex64)
        x = cg_solver(torch.zeros(a_dim, dtype=torch.complex64), y)
        x_ref = torch.solve(y, a_mat)[0]
        self.assertTrue(torch.allclose(x, x_ref, atol=1e-6))
//  [python]
def test_complex_valued_inverse(self):
        a_dim = 3
        a_mat = torch.tensor([[1, 2, 3], [2, 1, 2], [3, 2, 1]], dtype=torch.complex64)

        def a_op(x):
            return a_mat @ x

        cg_solver = ConjugateGradient(a_op, num_iter=100)
        y = torch.tensor([1, 2, 3], dtype=torch.complex64)
        x = cg_solver(torch.zeros(a_dim, dtype=torch.complex64), y)
        x_ref = torch.linalg.solve(a_mat, y)
        self.assertTrue(torch.allclose(x, x_ref, atol=1e-6))
//  [python]
def test_real_valued_inverse(self):
        """Test ConjugateGradient with real-valued input: when the input is real
        value, the output should be the inverse of the matrix."""
        a_dim = 3
        a_mat = torch.tensor([[1, 2, 3], [2, 1, 2], [3, 2, 1]], dtype=torch.float)

        def a_op(x):
            return a_mat @ x

        cg_solver = ConjugateGradient(a_op, num_iter=100)
        # define the measurement
        y = torch.tensor([1, 2, 3], dtype=torch.float)
        # solve for x
        x = cg_solver(torch.zeros(a_dim), y)
        x_ref = torch.solve(y, a_mat)[0]
        # assert torch.allclose(x, x_ref, atol=1e-6), 'CG solver failed to converge to reference solution'
        self.assertTrue(torch.allclose(x, x_ref, atol=1e-6))
//  [python]
def test_real_valued_inverse(self):
        """Test ConjugateGradient with real-valued input: when the input is real
        value, the output should be the inverse of the matrix."""
        a_dim = 3
        a_mat = torch.tensor([[1, 2, 3], [2, 1, 2], [3, 2, 1]], dtype=torch.float)

        def a_op(x):
            return a_mat @ x

        cg_solver = ConjugateGradient(a_op, num_iter=100)
        # define the measurement
        y = torch.tensor([1, 2, 3], dtype=torch.float)
        # solve for x
        x = cg_solver(torch.zeros(a_dim), y)
        x_ref = torch.linalg.solve(a_mat, y)
        # assert torch.allclose(x, x_ref, atol=1e-6), 'CG solver failed to converge to reference solution'
        self.assertTrue(torch.allclose(x, x_ref, atol=1e-6))
//  [python]
def forward(self, q: Tensor, v: Tensor,
                u: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
        """Evaluates multibody system dynamics terms at given state and input.

        Calculation is performed as a thin wrapper around
        :py:class:`LagrangianTerms` and :py:class:`ContactTerms`. For
        convenience, this function also returns the Delassus operator
        `D(q) = J(q)^T inv(M(q)) J(q)`.

        Args:
            q: (\*, n_q) configuration batch.
            v: (\*, n_v) velocity batch.
            u: (\*, n_u) input batch.

        Returns:
            (\*, 3 * n_collisions, 3 * n_collisions) Delassus operator D(q).
            (\*, n_v, n_v) mass matrix batch M(q).
            (\*, 3 * n_collisions, n_v) contact Jacobian J(q).
            (\*, n_collisions) signed distance phi(q).
            (\*, n_v) Contact-free acceleration inv(M(q)) * F(q).
        """
        M, non_contact_acceleration = self.lagrangian_terms(q, v, u)
        phi, J = self.contact_terms(q)

        delassus = pbmm(J, torch.solve(J.transpose(-1, -2), M)[0])
        return delassus, M, J, phi, non_contact_acceleration
//  [python]
def forward(self, q: Tensor, v: Tensor,
                u: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
        """Evaluates multibody system dynamics terms at given state and input.

        Calculation is performed as a thin wrapper around
        :py:class:`LagrangianTerms` and :py:class:`ContactTerms`. For
        convenience, this function also returns the Delassus operator
        `D(q) = J(q)^T inv(M(q)) J(q)`.

        Args:
            q: (\*, n_q) configuration batch.
            v: (\*, n_v) velocity batch.
            u: (\*, n_u) input batch.

        Returns:
            (\*, 3 * n_collisions, 3 * n_collisions) Delassus operator D(q).
            (\*, n_v, n_v) mass matrix batch M(q).
            (\*, 3 * n_collisions, n_v) contact Jacobian J(q).
            (\*, n_collisions) signed distance phi(q).
            (\*, n_v) Contact-free acceleration inv(M(q)) * F(q).
        """
        M, non_contact_acceleration = self.lagrangian_terms(q, v, u)
        phi, J = self.contact_terms(q)

        delassus = pbmm(J, torch.linalg.solve(M, J.transpose(-1, -2)))
        return delassus, M, J, phi, non_contact_acceleration
//  [python]
def train_closeform_MSE(self, features, y):
        inputs = features.cuda()
        labels = one_hot(y, num_labels=self.num_classes).cuda()
        A = torch.matmul(inputs.t(), inputs)  # .cpu()
        B = torch.matmul(inputs.t(), labels)  # .cpu()
        self.W = torch.matmul(torch.pinverse(A), B)  # .cuda()

        return self.W
//  [python]
def train_closeform_MSE(self, features, y):
        inputs = features.cuda()
        labels = one_hot(y, num_labels=self.num_classes).cuda()
        A = torch.matmul(inputs.t(), inputs)  # .cpu()
        B = torch.matmul(inputs.t(), labels)  # .cpu()
        # self.W = torch.matmul(torch.pinverse(A), B)  # .cuda()
        try:
            self.W = torch.linalg.solve(A, B)
        except:
            self.W = torch.matmul(torch.pinverse(A), B)

        return self.W
//  [python]
def test_caching_parametrization_with_transfer_parametrizations_and_params(self):
        r"""Test that transferring parametrizations doesn't cause issues with caching"""
        class Skew(nn.Module):
            def forward(self, X):
                X = X.tril(-1)
                return X - X.T

        class Orthogonal(nn.Module):
            def forward(self, X):
                Id = torch.eye(X.size(0), device=X.device)
                return torch.solve(Id - X, Id + X)[0]

        model = nn.Linear(5, 5)
        parametrize.register_parametrization(model, "weight", Skew())
        parametrize.register_parametrization(model, "weight", Orthogonal())

        to_model = nn.Linear(5, 5)
        parametrize.transfer_parametrizations_and_params(model, to_model)

        with parametrize.cached():
            X = model.weight
            Y = model.weight
            self.assertEqual(id(X), id(Y))

            A = to_model.weight
            B = to_model.weight
            self.assertEqual(id(A), id(B))

            # test that the results are distinct objects for each module
            self.assertNotEqual(id(A), id(X))
//  [python]
def test_caching_parametrization_with_transfer_parametrizations_and_params(self):
        r"""Test that transferring parametrizations doesn't cause issues with caching"""
        class Skew(nn.Module):
            def forward(self, X):
                X = X.tril(-1)
                return X - X.T

        class Orthogonal(nn.Module):
            def forward(self, X):
                Id = torch.eye(X.size(0), device=X.device)
                return torch.linalg.solve(Id + X, Id - X)

        model = nn.Linear(5, 5)
        parametrize.register_parametrization(model, "weight", Skew())
        parametrize.register_parametrization(model, "weight", Orthogonal())

        to_model = nn.Linear(5, 5)
        parametrize.transfer_parametrizations_and_params(model, to_model)

        with parametrize.cached():
            X = model.weight
            Y = model.weight
            self.assertEqual(id(X), id(Y))

            A = to_model.weight
            B = to_model.weight
            self.assertEqual(id(A), id(B))

            # test that the results are distinct objects for each module
            self.assertNotEqual(id(A), id(X))
//  [python]
def solve_weights(t_i, gamma_i, t_r, gamma_r):
    rv = lambda x: x.double().reshape(-1, 1)
    cv = lambda x: x.double().reshape(1, -1)
    A = p_dot_p(rv(t_i), rv(gamma_i), cv(t_i), cv(gamma_i))
    b = p_dot_p(rv(t_i), rv(gamma_i), cv(t_r), cv(gamma_r))
    return torch.solve(b, A)[0]
//  [python]
def solve_weights(t_i, gamma_i, t_r, gamma_r):
    rv = lambda x: x.double().reshape(-1, 1)
    cv = lambda x: x.double().reshape(1, -1)
    A = p_dot_p(rv(t_i), rv(gamma_i), cv(t_i), cv(gamma_i))
    b = p_dot_p(rv(t_i), rv(gamma_i), cv(t_r), cv(gamma_r))
    return torch.linalg.solve(A, b)
//  [python]
def get_torch_solve(
        self,
        A, b,
        method,
    ):
        if method == 'inverse':
            return torch.solve(b, A)[0]
        elif method == 'cholesky':
            l = torch.cholesky(A)
            z = torch.triangular_solve(b, l, transpose=False, upper=False)[0]
            return torch.triangular_solve(z, l, transpose=True, upper=False)[0]
        else:
            raise NotImplementedError
//  [python]
def get_torch_solve(
        self,
        A, b,
        method,
    ):
        if method == 'inverse':
            return torch.linalg.solve(A, b)
        elif method == 'cholesky':
            l = torch.linalg.cholesky(A)
            z = torch.linalg.solve_triangular(l, b, upper=False)
            return torch.linalg.solve_triangular(l.mT, z, upper=False)
            # z = torch.triangular_solve(b, l, transpose=False, upper=False)[0]
            # return torch.triangular_solve(z, l, transpose=True, upper=False)[0]
        else:
            raise NotImplementedError
//  [python]
def score_parameters(self) -> List[Tensor]:
        """
        :return: List of Tensors the same shapes as the given Parameters where
            each Parameter's elements are scored based on the blockwise OBS
        """
        scores = [None] * len(self._params)
        block_finv_w = [None] * len(self._params)

        if self._is_main_proc:
            for i, finv in enumerate(self._finvs):
                if self._mask_type == "unstructured" or "uvp" in self._mask_type:
                    scores[i] = (
                        (self._params[i].data.reshape(-1) ** 2).to(self._devices[i])
                        / (2.0 * finv.diag() + self._eps)
                    ).reshape(self._params[i].shape)
                else:  # self._mask_type == "block4":
                    block_w = (
                        self._params[i].data.reshape(-1, 4).to(finv.dev)
                    )  # (d/Q, Q)
                    block_finv = (
                        torch.cat(
                            [
                                finv.f_inv[:, i : i + 4, i : i + 4]
                                for i in range(0, finv.B, 4)
                            ],
                            dim=1,
                        )
                        .reshape((finv.d // finv.B, finv.B // 4, 4, 4))
                        .reshape((finv.d // 4, 4, 4))
                    )  # (Q, d/Q, Q) -> (d/Q, Q, Q)
                    block_finv_w[i] = torch.solve(
                        block_w,
                        block_finv,
                    )[0]  # (d/Q, Q)
                    score = 0.5 * torch.einsum(
                        "bi,bi->b", block_w, block_finv_w[i]
                    )  # d/Q
                    scores[i] = (
                        score.unsqueeze(1)
                        .expand(-1, 4)
                        .reshape(self._params[i].data.shape)
                    )

            # make sure pruned ones will stay pruned
            for i, score in enumerate(scores):
                score[self._masks[i] == 0] = float("-inf")

        self._broadcast_list_from_main(scores)
        if self._mask_type == "block4":
            self._broadcast_list_from_main(block_finv_w)
            self._block_finv_w = block_finv_w  # cache for OBS weight update

        return scores
//  [python]
def score_parameters(self) -> List[Tensor]:
        """
        :return: List of Tensors the same shapes as the given Parameters where
            each Parameter's elements are scored based on the blockwise OBS
        """
        scores = [None] * len(self._params)
        block_finv_w = [None] * len(self._params)

        if self._is_main_proc:
            for i, finv in enumerate(self._finvs):
                if self._mask_type == "unstructured" or "uvp" in self._mask_type:
                    scores[i] = (
                        (self._params[i].data.reshape(-1) ** 2).to(self._devices[i])
                        / (2.0 * finv.diag() + self._eps)
                    ).reshape(self._params[i].shape)
                else:  # self._mask_type == "block4":
                    block_w = (
                        self._params[i].data.reshape(-1, 4).to(finv.dev)
                    )  # (d/Q, Q)
                    block_finv = (
                        torch.cat(
                            [
                                finv.f_inv[:, i : i + 4, i : i + 4]
                                for i in range(0, finv.B, 4)
                            ],
                            dim=1,
                        )
                        .reshape((finv.d // finv.B, finv.B // 4, 4, 4))
                        .reshape((finv.d // 4, 4, 4))
                    )  # (Q, d/Q, Q) -> (d/Q, Q, Q)
                    block_finv_w[i] = torch.linalg.solve(
                        block_finv,
                        block_w,
                    )  # (d/Q, Q)
                    score = 0.5 * torch.einsum(
                        "bi,bi->b", block_w, block_finv_w[i]
                    )  # d/Q
                    scores[i] = (
                        score.unsqueeze(1)
                        .expand(-1, 4)
                        .reshape(self._params[i].data.shape)
                    )

            # make sure pruned ones will stay pruned
            for i, score in enumerate(scores):
                score[self._masks[i] == 0] = float("-inf")

        self._broadcast_list_from_main(scores)
        if self._mask_type == "block4":
            self._broadcast_list_from_main(block_finv_w)
            self._block_finv_w = block_finv_w  # cache for OBS weight update

        return scores
//  [python]
def fit(self, X: torch.Tensor, Y: torch.Tensor):
        r"""Learn f that matches Y given X

        Args:
            X (Tensor): Control point at source space (X_c)
                Shape: (n_c, d_s)
            Y (Tensor): Control point in the target space (X_t)
                Shape: (n_c, d_t)

        Returns:
            ThinPlateSpline: self
        """
        X = X.to(self.device)
        Y = Y.to(self.device)
        X = _ensure_2d(X)
        Y = _ensure_2d(Y)
        assert X.shape[0] == Y.shape[0]

        n_c, d_s = X.shape
        self.control_points = X

        phi = self._radial_distance(X)

        # Build the linear system AP = Y
        X_p = torch.hstack([torch.ones((n_c, 1), device=self.device), X])

        A = torch.vstack(
            [
                torch.hstack(
                    [
                        phi + self.alpha * torch.eye(n_c, device=self.device),
                        X_p,
                    ]
                ),
                torch.hstack(
                    [
                        X_p.T,
                        torch.zeros((d_s + 1, d_s + 1), device=self.device),
                    ]
                ),
            ]
        )

        Y = torch.vstack([Y, torch.zeros((d_s + 1, Y.shape[1]), device=self.device)])

        self.parameters = torch.solve(Y, A)[1]  # pylint: disable=not-callable
        self._fitted = True

        return self
//  [python]
def fit(self, X: torch.Tensor, Y: torch.Tensor):
        r"""Learn f that matches Y given X

        Args:
            X (Tensor): Control point at source space (X_c)
                Shape: (n_c, d_s)
            Y (Tensor): Control point in the target space (X_t)
                Shape: (n_c, d_t)

        Returns:
            ThinPlateSpline: self
        """
        X = X.to(self.device)
        Y = Y.to(self.device)
        X = _ensure_2d(X)
        Y = _ensure_2d(Y)
        assert X.shape[0] == Y.shape[0]

        n_c, d_s = X.shape
        self.control_points = X

        phi = self._radial_distance(X)

        # Build the linear system AP = Y
        X_p = torch.hstack([torch.ones((n_c, 1), device=self.device), X])

        A = torch.vstack(
            [
                torch.hstack(
                    [
                        phi + self.alpha * torch.eye(n_c, device=self.device),
                        X_p,
                    ]
                ),
                torch.hstack(
                    [
                        X_p.T,
                        torch.zeros((d_s + 1, d_s + 1), device=self.device),
                    ]
                ),
            ]
        )

        Y = torch.vstack([Y, torch.zeros((d_s + 1, Y.shape[1]), device=self.device)])

        self.parameters = torch.linalg.solve(A, Y)  # pylint: disable=not-callable
        self._fitted = True

        return self
//  [python]
def quadratic_flow(F0n1: torch.Tensor, F01: torch.Tensor, gamma: float, tau: float) -> torch.Tensor:
    """solve the quadratic motion matrix and predict the correction feild.
    Args:
        F0n1: torch.Tensor (b,h,w,2), flow 0 -> -1
        F01: torch.Tensor (b,h,w,2), flow 0 -> 1
        gamma (float): the readout reatio
        tau (float): the timestamp warping to
    Returns:
        torch.Tensor: the correction feild to tau.
    """
    h, w = F0n1.shape[1:3]
    t0n1 = -1 + gamma / h * F0n1[:, :, :, 1]  # * (b, h, w)
    t01 = 1 + gamma / h * F01[:, :, :, 1]  # * (b, h, w)

    # solve the quadratic motion matrix
    A = rearrange(
        torch.stack([t0n1, 0.5 * t0n1**2, t01, 0.5 * t01**2], dim=-1),
        "b h w (m n) -> b h w m n",
        m=2,
        n=2,
    )  # * (b, h, w, 2, 2)

    B = torch.stack([F0n1, F01], dim=-2)  # * (b, h, w, 2, 2)
    M = torch.solve(B, A)[0]  # * (b, h, w, 2, 2)

    # predict the correction feild
    grid_y, _ = torch.meshgrid(
        torch.arange(0, h, device=F0n1.device, requires_grad=False),
        torch.arange(0, w, device=F0n1.device, requires_grad=False),
    )
    t0tau = tau - gamma / h * grid_y  # * (h, w)

    Atau = rearrange(torch.stack([t0tau, 0.5 * t0tau**2], dim=-1), "h w m -> h w 1 m")  # * (h, w, 1, 2)
    F0tau = rearrange(Atau @ M, "b h w 1 n -> b h w n")  # * (b, h, w, 2)

    return F0tau
//  [python]
def quadratic_flow(F0n1: torch.Tensor, F01: torch.Tensor, gamma: float, tau: float) -> torch.Tensor:
    """solve the quadratic motion matrix and predict the correction feild.
    Args:
        F0n1: torch.Tensor (b,h,w,2), flow 0 -> -1
        F01: torch.Tensor (b,h,w,2), flow 0 -> 1
        gamma (float): the readout reatio
        tau (float): the timestamp warping to
    Returns:
        torch.Tensor: the correction feild to tau.
    """
    h, w = F0n1.shape[1:3]
    t0n1 = -1 + gamma / h * F0n1[:, :, :, 1]  # * (b, h, w)
    t01 = 1 + gamma / h * F01[:, :, :, 1]  # * (b, h, w)

    # solve the quadratic motion matrix
    A = rearrange(
        torch.stack([t0n1, 0.5 * t0n1**2, t01, 0.5 * t01**2], dim=-1),
        "b h w (m n) -> b h w m n",
        m=2,
        n=2,
    )  # * (b, h, w, 2, 2)

    B = torch.stack([F0n1, F01], dim=-2)  # * (b, h, w, 2, 2)
    M = torch.linalg.solve(A, B)  # * (b, h, w, 2, 2)

    # predict the correction feild
    grid_y, _ = torch.meshgrid(
        torch.arange(0, h, device=F0n1.device, requires_grad=False),
        torch.arange(0, w, device=F0n1.device, requires_grad=False),
    )
    t0tau = tau - gamma / h * grid_y  # * (h, w)

    Atau = rearrange(torch.stack([t0tau, 0.5 * t0tau**2], dim=-1), "h w m -> h w 1 m")  # * (h, w, 1, 2)
    F0tau = rearrange(Atau @ M, "b h w 1 n -> b h w n")  # * (b, h, w, 2)

    return F0tau
//  [python]
def direct_newton(env, cmd, reg_ctrl, variant='gmres_scipy', time_solve_only=False):
    tic0 = time.time()
    horizon, dim_ctrl = cmd.shape

    # Compute step from a given command
    cmd_flat = deepcopy(cmd.data)
    cmd_flat = cmd_flat.view(-1)
    cmd_flat.requires_grad = True
    cmd_aux = cmd_flat.view(horizon, dim_ctrl)

    _, costs = env.forward(cmd_aux, approx='linquad')
    total_cost = sum(costs)

    # Get gradient, hessian, and make a newton step to get the solution
    grad = torch.autograd.grad(total_cost, cmd_flat, create_graph=True)[0]
    # add regularization in the hessian
    if variant == 'torch_solve':
      hess = auto_multi_grad(grad, cmd_flat) + reg_ctrl * torch.eye(dim_ctrl * horizon)
      tic1 = time.time()
      cmd_opt_newton = - torch.solve(grad.unsqueeze(-1), hess)[0].view(-1)
    elif variant == 'own_cg':
      hvp = make_hvp(grad, cmd_flat)
      hvp_reg = lambda x: hvp(x) + reg_ctrl*x
      tic1 = time.time()
      cmd_opt_newton = own_cg(hvp_reg, grad)
    elif variant == 'gmres_scipy':
      hvp = make_numpy_hvp(grad, cmd_flat)
      hvp_reg = lambda x: hvp(x) + reg_ctrl*x
      hvp_reg = LinearOperator((horizon*dim_ctrl, horizon*dim_ctrl), hvp_reg)
      target = deepcopy(grad.data).numpy()
      tic1 = time.time()
      cmd_opt_newton = cg(hvp_reg, target)[0]
      cmd_opt_newton = torch.from_numpy(cmd_opt_newton)
    else:
        raise NotImplementedError


    cmd_opt_newton = cmd_opt_newton.view(horizon, dim_ctrl)
    time_newton = time.time() - tic1 if time_solve_only else time.time() - tic0
    return cmd_opt_newton, time_newton