jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _get_perspective_coeffs(
        startpoints: List[List[int]], endpoints: List[List[int]]
) -> List[float]:
    """Helper function to get the coefficients (a, b, c, d, e, f, g, h) for the perspective transforms.

    In Perspective Transform each pixel (x, y) in the original image gets transformed as,
     (x, y) -> ( (ax + by + c) / (gx + hy + 1), (dx + ey + f) / (gx + hy + 1) )

    Args:
        startpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the original image.
        endpoints (list of list of ints): List containing four lists of two integers corresponding to four corners
            ``[top-left, top-right, bottom-right, bottom-left]`` of the transformed image.

    Returns:
        octuple (a, b, c, d, e, f, g, h) for transforming each pixel.
    """
    a_matrix = torch.zeros(2 * len(startpoints), 8, dtype=torch.float)

    for i, (p1, p2) in enumerate(zip(endpoints, startpoints)):
        a_matrix[2 * i, :] = torch.tensor([p1[0], p1[1], 1, 0, 0, 0, -p2[0] * p1[0], -p2[0] * p1[1]])
        a_matrix[2 * i + 1, :] = torch.tensor([0, 0, 0, p1[0], p1[1], 1, -p2[1] * p1[0], -p2[1] * p1[1]])

    b_matrix = torch.tensor(startpoints, dtype=torch.float).view(8)
    res = torch.linalg.lstsq(a_matrix, b_matrix, driver='gels').solution

    output: List[float] = res.tolist()
    return output
//  [python]
def forward(self, A: Tensor, b: Tensor) -> Tensor:
        '''
        Args:
            A (Tensor): the input batched tensor.
            b (Tensor): the batched tensor on the right hand side.

        Return:
            Tensor: the solved batched tensor.
        '''
        self.out = torch.lstsq(b, A, rcond=self.rcond)
        assert not torch.any(torch.isnan(self.out[0])), \
            'Linear Solver Failed Using LSTSQ. Using PINV() instead'
        return self.out[0]
//  [python]
def forward(self, A: Tensor, b: Tensor) -> Tensor:
        '''
        Args:
            A (Tensor): the input batched tensor.
            b (Tensor): the batched tensor on the right hand side.

        Return:
            Tensor: the solved batched tensor.
        '''
        self.out = lstsq(A, b, rcond=self.rcond, driver=self.driver)
        assert not torch.any(torch.isnan(self.out.solution)), \
            'Linear Solver Failed Using LSTSQ. Using PINV() instead'
        return self.out.solution
//  [python]
def fine_ture(self, x, y):
        # optimization using extreme learning machine
        m, s = x.mean(), x.std()
        x = (x - m) / s
        y = (y - m) / s
        y = y.unsqueeze(dim=0).to(torch.float) if y.dim() < 2 else y.to(torch.float)
        hid_out = self.act(self.inp2hid(x.to(torch.float)))
        self.hid2out.state_dict()['weight'] = torch.solve(torch.matmul(hid_out.t(), y), torch.matmul(hid_out.t(), hid_out)).solution
//  [python]
def fine_ture(self, x, y):
        # optimization using extreme learning machine
        m, s = x.mean(), x.std()
        x = (x - m) / s
        y = (y - m) / s
        y = y.unsqueeze(dim=0).to(torch.float) if y.dim() < 2 else y.to(torch.float)
        hid_out = self.act(self.inp2hid(x.to(torch.float)))
        self.hid2out.state_dict()['weight'] = torch.linalg.lstsq(hid_out, y.permute(1, 0)).solution
//  [python]
def forward(self):
        camera = self.get_normalized_camera()
        acamera = self.get_autoencoded_camera()
        
        # Loss terms ###########################################################                
        autoencoder_error = torch.acos(
            torch.clamp(                
                torch.dot(camera.flatten(), acamera.flatten())/(camera.norm()*acamera.norm() + 1e-7),
                min = -1.0 + 1e-7,
                max = 1.0 - 1e-7,
            )            
        )

        # C-matrices
        predicted_cmatrices = [torch.solve(L.T*camera, L.T*self.observer)[0]
                               for L in self.cilluminants]
        cerror = sum(
                torch.acos(
                    torch.clamp(
                        (torch.dot(pred_cmat.flatten(), cmat.flatten()))/(pred_cmat.norm()*cmat.norm() + 1e-7),
                        min = -1.0 + 1e-7,
                        max = 1.0 - 1e-7,
                    ),                
                
            )
            for cmat, pred_cmat in zip(self.cmatrices, predicted_cmatrices)
        )

        norm_error = (self.camera.norm() - 1.0)**2

        loss = sum([
            self.config["autoencoder"]*autoencoder_error,
            self.config["cmatrices"]*cerror,
        ])        
               
        return loss
//  [python]
def forward(self):
        camera = self.get_normalized_camera()
        acamera = self.get_autoencoded_camera()
        
        # Loss terms ###########################################################                
        autoencoder_error = torch.acos(
            torch.clamp(                
                torch.dot(camera.flatten(), acamera.flatten())/(camera.norm()*acamera.norm() + 1e-7),
                min = -1.0 + 1e-7,
                max = 1.0 - 1e-7,
            )            
        )

        # C-matrices
        predicted_cmatrices = [torch.linalg.lstsq(L.T*camera, L.T*self.observer).solution
                               for L in self.cilluminants]
        cerror = sum(
                torch.acos(
                    torch.clamp(
                        (torch.dot(pred_cmat.flatten(), cmat.flatten()))/(pred_cmat.norm()*cmat.norm() + 1e-7),
                        min = -1.0 + 1e-7,
                        max = 1.0 - 1e-7,
                    ),                
                
            )
            for cmat, pred_cmat in zip(self.cmatrices, predicted_cmatrices)
        )

        norm_error = (self.camera.norm() - 1.0)**2

        loss = sum([
            self.config["autoencoder"]*autoencoder_error,
            self.config["cmatrices"]*cerror,
        ])        
               
        return loss
//  [python]
def inverse(self, output):
        b, c, h, w = output.shape
        A = self.weight[..., 0,0] # outch, inch
        B = output.permute([1,0,2,3]).reshape(c, -1) # outch, bhw
        X = torch.solve(B, A.T)  # AX=B
        return X.solution.reshape(-1, b, h, w).permute([1, 0, 2, 3])
//  [python]
def inverse(self, output):
        b, c, h, w = output.shape
        A = self.weight[..., 0,0] # outch, inch
        B = output.permute([1,0,2,3]).reshape(c, -1) # outch, bhw
        X = torch.linalg.lstsq(A, B)  # AX=B
        return X.solution.reshape(-1, b, h, w).permute([1, 0, 2, 3])
//  [python]
def get_ray_intersections(ray1, ray2):
    """
    calculates points on ray1 and ray2 where both rays are closest to another
    :param ray1: torch Tensor [orgx, orgy, orgz, dirx, diry, dirz]
    :param ray2: torch Tensor [orgx, orgy, orgz, dirx, diry, dirz]
    :return:
    """

    B = (ray2[:3] - ray1[:3]).unsqueeze(1)
    A = torch.stack((ray1[3:], -ray2[3:]), dim=-1)

    t1t2 = torch.solve(B, A).solution
    t1t2 = t1t2.flatten()

    x1 = ray1[:3] + ray1[3:] * t1t2[0]
    x2 = ray2[:3] + ray2[3:] * t1t2[1]

    return x1, x2
//  [python]
def get_ray_intersections(ray1, ray2):
    """
    calculates points on ray1 and ray2 where both rays are closest to another
    :param ray1: torch Tensor [orgx, orgy, orgz, dirx, diry, dirz]
    :param ray2: torch Tensor [orgx, orgy, orgz, dirx, diry, dirz]
    :return:
    """

    B = (ray2[:3] - ray1[:3]).unsqueeze(1)
    A = torch.stack((ray1[3:], -ray2[3:]), dim=-1)

    t1t2 = torch.linalg.lstsq(A, B).solution
    t1t2 = t1t2.flatten()

    x1 = ray1[:3] + ray1[3:] * t1t2[0]
    x2 = ray2[:3] + ray2[3:] * t1t2[1]

    return x1, x2
//  [python]
def solve(self, vs, regul=1e-8, use_pi=True):
        vs_dict = vs.get_dict_representation()
        out_dict = dict()
        for layer_id, layer in self.generator.layer_collection.layers.items():
            vw = vs_dict[layer_id][0]
            sw = vw.size()
            v = vw.view(sw[0], -1)
            if layer.bias is not None:
                v = torch.cat([v, vs_dict[layer_id][1].unsqueeze(1)], dim=1)
            a, g = self.data[layer_id]
            if use_pi:
                pi = (torch.trace(a) / torch.trace(g) * g.size(0) / a.size(0)) ** 0.5
            else:
                pi = 1
            a_reg = a + regul**0.5 * pi * torch.eye(a.size(0), device=g.device)
            g_reg = g + regul**0.5 / pi * torch.eye(g.size(0), device=g.device)

            solve_g, _ = torch.gesv(v, g_reg)
            solve_a, _ = torch.gesv(solve_g.t(), a_reg)
            solve_a = solve_a.t()
            if layer.bias is None:
                solve_tuple = (solve_a.view(*sw),)
            else:
                solve_tuple = (
                    solve_a[:, :-1].contiguous().view(*sw),
                    solve_a[:, -1].contiguous(),
                )
            out_dict[layer_id] = solve_tuple
        return PVector(layer_collection=vs.layer_collection, dict_repr=out_dict)
//  [python]
def solve(self, vs, regul=1e-8, use_pi=True):
        vs_dict = vs.get_dict_representation()
        out_dict = dict()
        for layer_id, layer in self.generator.layer_collection.layers.items():
            vw = vs_dict[layer_id][0]
            sw = vw.size()
            v = vw.view(sw[0], -1)
            if layer.bias is not None:
                v = torch.cat([v, vs_dict[layer_id][1].unsqueeze(1)], dim=1)
            a, g = self.data[layer_id]
            if use_pi:
                pi = (torch.trace(a) / torch.trace(g) * g.size(0) / a.size(0)) ** 0.5
            else:
                pi = 1
            a_reg = a + regul**0.5 * pi * torch.eye(a.size(0), device=g.device)
            g_reg = g + regul**0.5 / pi * torch.eye(g.size(0), device=g.device)

            solve_g, _, _, _ = torch.linalg.lstsq(g_reg, v)
            solve_a, _, _, _ = torch.linalg.lstsq(a_reg, solve_g.t())
            solve_a = solve_a.t()
            if layer.bias is None:
                solve_tuple = (solve_a.view(*sw),)
            else:
                solve_tuple = (
                    solve_a[:, :-1].contiguous().view(*sw),
                    solve_a[:, -1].contiguous(),
                )
            out_dict[layer_id] = solve_tuple
        return PVector(layer_collection=vs.layer_collection, dict_repr=out_dict)
//  [python]
def proj_range(self, x: Tensor):
        return (
            x
            + torch.lstsq(
                self.A,
                self.y - self.A @ x,
            )[0]
        )
//  [python]
def proj_range(self, x: Tensor):
        return (
            x
            + torch.linalg.lstsq(
                self.A,
                self.y - self.A @ x,
            )[0]
        )
//  [python]
def SINDy(Theta, dXdt, lamb=0.05):
    d = dXdt.shape[1]
    Xi = torch.lstsq(Theta, dXdt)[0] # Initial guess: Least-squares
    
    for _ in range(10): #which parameter should we be tuning here for RRR comp
        smallinds = torch.abs(Xi) < lamb # Find small coefficients
        Xi[smallinds] = 0             # and threshold
        for ind in range(d):          # n is state dimension
            biginds = smallinds[:, ind] == 0
            # Regress dynamics onto remaining terms to find sparse Xi
            Xi[biginds, ind] = np.linalg.lstsq(Theta[:, biginds], dXdt[:, ind], rcond=None)[0]

    L = Xi
    return L
//  [python]
def SINDy(Theta, dXdt, lamb=0.05):
    d = dXdt.shape[1]
    Xi = torch.linalg.lstsq(Theta, dXdt, rcond=None)[0] # Initial guess: Least-squares
    
    for _ in range(10): #which parameter should we be tuning here for RRR comp
        smallinds = torch.abs(Xi) < lamb # Find small coefficients
        Xi[smallinds] = 0             # and threshold
        for ind in range(d):          # n is state dimension
            biginds = smallinds[:, ind] == 0
            # Regress dynamics onto remaining terms to find sparse Xi
            Xi[biginds, ind] = np.linalg.lstsq(Theta[:, biginds], dXdt[:, ind], rcond=None)[0]

    L = Xi
    return L
//  [python]
def closed_form_april(model, shared_data, solver="gelsd"):
        """Run inversions on CPU in double precision. (gelsd only implemented for CPU)"""
        # recover patch embeddings first, (APRIL paper)
        qkv_w = model.model.blocks[0].attn.qkv.weight.detach().double().cpu()
        q_w, k_w, v_w = qkv_w.reshape(3, -1, qkv_w.shape[-1]).unbind()
        qkv_g = shared_data["gradients"][4].detach().double().cpu()
        assert qkv_w.shape == qkv_g.shape
        q_g, k_g, v_g = qkv_g.reshape(3, -1, qkv_g.shape[-1]).unbind()
        A = shared_data["gradients"][1].detach().squeeze().double().cpu()
        pos_embed = model.model.pos_embed.detach().squeeze().double().cpu()

        b = (q_w.T @ q_g + k_w.T @ k_g + v_w.T @ v_g).double().cpu()
        log.info(f"Attention Inversion:  ||A||={A.norm()}, ||b||={b.norm()}")
        z = torch.solve(A.T @ b, A.T @ A).solution
        z -= pos_embed
        z = z[1:]

        # recover img
        em_w = model.model.patch_embed.proj.weight.detach().double().cpu()
        em_w = em_w.reshape((em_w.shape[0], -1))
        em_b = model.model.patch_embed.proj.bias.detach().double().cpu()

        x = z - em_b
        log.info(f"Embedding Inversion:  ||A||={em_w.norm()}, ||b||={x.norm()}")
        x = torch.solve(em_w.T @ x.T, em_w.T @ em_w).solution
        x = x.reshape((3, -1, x.shape[-1]))
        x = x.transpose(1, 2)
        return x
//  [python]
def closed_form_april(model, shared_data, solver="gelsd"):
        """Run inversions on CPU in double precision. (gelsd only implemented for CPU)"""
        # recover patch embeddings first, (APRIL paper)
        qkv_w = model.model.blocks[0].attn.qkv.weight.detach().double().cpu()
        q_w, k_w, v_w = qkv_w.reshape(3, -1, qkv_w.shape[-1]).unbind()
        qkv_g = shared_data["gradients"][4].detach().double().cpu()
        assert qkv_w.shape == qkv_g.shape
        q_g, k_g, v_g = qkv_g.reshape(3, -1, qkv_g.shape[-1]).unbind()
        A = shared_data["gradients"][1].detach().squeeze().double().cpu()
        pos_embed = model.model.pos_embed.detach().squeeze().double().cpu()

        b = (q_w.T @ q_g + k_w.T @ k_g + v_w.T @ v_g).double().cpu()
        log.info(f"Attention Inversion:  ||A||={A.norm()}, ||b||={b.norm()}")
        z = torch.linalg.lstsq(A.T, b, driver=solver, rcond=None).solution
        z -= pos_embed
        z = z[1:]

        # recover img
        em_w = model.model.patch_embed.proj.weight.detach().double().cpu()
        em_w = em_w.reshape((em_w.shape[0], -1))
        em_b = model.model.patch_embed.proj.bias.detach().double().cpu()

        x = z - em_b
        log.info(f"Embedding Inversion:  ||A||={em_w.norm()}, ||b||={x.norm()}")
        x = torch.linalg.lstsq(em_w, x.T, driver=solver, rcond=None).solution
        x = x.reshape((3, -1, x.shape[-1]))
        x = x.transpose(1, 2)
        return x
//  [python]
def barycentric_interpolation(
    rotations: Tensor,
    grid: Tensor,
    signal: Tensor,
    *,
    dist_fn=geodesic_distance,
    eps: float = 1e-3,
) -> Tensor:
    """
    Performs linear interpolation of rotations using barycentric coordinates.

    Arguments:
        - rotations: Tensor of shape `(N, K, 4)` containing quaternions.
        - grid: Tensor of shape `(N, L, 4)` containing quaternions
                defining a grid on SO3.
        - signal: Tensor of shape `(N, L, S)` containing the signal
                  corresponding to `grid`.
        - dist_fn: Distance function used to calculate distance between
                   rotations.
        - eps: Float for preventing numerical instabilities.

    Returns:
        - Tensor of shape `(N, K, S)`.
    """
    N, _, S = signal.shape
    _, H, D = rotations.shape

    dists = dist_fn(rotations[..., None, :], grid[..., None, :, :])
    dists_k, idx = torch.topk(dists, k=3, largest=False)

    simplices = (
        grid[..., None, :]
        .expand(-1, -1, D, -1)
        .gather(1, idx[..., None].expand(-1, -1, -1, D))
    )

    # to make sure all quaternions live on same shard
    simplices *= 2 * (simplices @ rotations[..., None] > 0) - 1

    bcc = torch.solve(rotations[..., None], simplices.transpose(-1, -2))[0].view(
        -1, 3
    )

    mask = dists_k[:, :, 0].view(-1) <= eps
    bcc[mask, 0], bcc[mask, 1], bcc[mask, 2] = 1.0, 0.0, 0.0
    bcc[bcc < 0] = 0

    bcc = bcc.view(N, H, 3)
    bcc /= bcc.sum(-1, keepdim=True)

    signal = (
        signal[..., None, :]
        .expand(-1, -1, 3, -1)
        .gather(1, idx[..., None].expand(-1, -1, -1, S))
    )

    return torch.sum(bcc[..., None] * signal, axis=-2)
//  [python]
def barycentric_interpolation(
    rotations: Tensor,
    grid: Tensor,
    signal: Tensor,
    *,
    dist_fn=geodesic_distance,
    eps: float = 1e-3,
) -> Tensor:
    """
    Performs linear interpolation of rotations using barycentric coordinates.

    Arguments:
        - rotations: Tensor of shape `(N, K, 4)` containing quaternions.
        - grid: Tensor of shape `(N, L, 4)` containing quaternions
                defining a grid on SO3.
        - signal: Tensor of shape `(N, L, S)` containing the signal
                  corresponding to `grid`.
        - dist_fn: Distance function used to calculate distance between
                   rotations.
        - eps: Float for preventing numerical instabilities.

    Returns:
        - Tensor of shape `(N, K, S)`.
    """
    N, _, S = signal.shape
    _, H, D = rotations.shape

    dists = dist_fn(rotations[..., None, :], grid[..., None, :, :])
    dists_k, idx = torch.topk(dists, k=3, largest=False)

    simplices = (
        grid[..., None, :]
        .expand(-1, -1, D, -1)
        .gather(1, idx[..., None].expand(-1, -1, -1, D))
    )

    # to make sure all quaternions live on same shard
    simplices *= 2 * (simplices @ rotations[..., None] > 0) - 1

    bcc = torch.linalg.lstsq(simplices.transpose(-1, -2), rotations[..., None])[0].view(
        -1, 3
    )

    mask = dists_k[:, :, 0].view(-1) <= eps
    bcc[mask, 0], bcc[mask, 1], bcc[mask, 2] = 1.0, 0.0, 0.0
    bcc[bcc < 0] = 0

    bcc = bcc.view(N, H, 3)
    bcc /= bcc.sum(-1, keepdim=True)

    signal = (
        signal[..., None, :]
        .expand(-1, -1, 3, -1)
        .gather(1, idx[..., None].expand(-1, -1, -1, S))
    )

    return torch.sum(bcc[..., None] * signal, axis=-2)
//  [python]
def lsq_gna(
        p: torch.Tensor,
        function: Callable, 
        jac_function: Callable = None,
        args: Union[Tuple, List] = (),
        ftol: float = 1e-8,
        ptol: float = 1e-8,
        gtol: float = 1e-8,
        l: float = 1.,
        max_iter: int = 100,
    ) -> List[torch.Tensor]:
    """
    Gauss-Newton implementation for least-squares fitting of non-linear functions
    
    :param p: initial value(s)
    :param function: user-provided function which takes p (and additional arguments) as input
    :param jac_fun: user-provided Jacobian function which takes p (and additional arguments) as input
    :param args: optional arguments passed to function
    :param ftol: relative change in cost function as stop condition
    :param ptol: relative change in independant variables as stop condition
    :param gtol: maximum gradient tolerance as stop condition
    :param l: step size damping parameter
    :param max_iter: maximum number of iterations
    :return: list of results
    """

    if len(args) > 0:
        # pass optional arguments to function
        fun = lambda p: function(p, *args)
    else:
        fun = function

    if jac_function is None:
        # use numerical Jacobian if analytical is not provided
        jac_fun = lambda p: jacobian_approx_t(p, f=fun)
    else:
        jac_fun = lambda p: jac_function(p, *args)

    f = fun(p)
    j = jac_fun(p)
    g = torch.matmul(j.T, f)
    H = torch.matmul(j.T, j)
    p_list = [p]
    while len(p_list) < max_iter:
        h = -l*torch.solve(g,H)[0]
        p = p + h
        p_list.append(p.clone())
        f_prev = f.clone()
        f = fun(p)
        j = jac_fun(p)
        g = torch.matmul(j.T, f)
        H = torch.matmul(j.T, j)

        # stop conditions
        gcon = max(abs(g)) < gtol
        pcon = (h**2).sum()**.5 < ptol*(ptol + (p**2).sum()**.5)
        fcon = ((f_prev-f)**2).sum() < ((ftol*f)**2).sum()
        if gcon or pcon or fcon:
            break

    return p_list
//  [python]
def lsq_gna(
        p: torch.Tensor,
        function: Callable, 
        jac_function: Callable = None,
        args: Union[Tuple, List] = (),
        ftol: float = 1e-8,
        ptol: float = 1e-8,
        gtol: float = 1e-8,
        l: float = 1.,
        max_iter: int = 100,
    ) -> List[torch.Tensor]:
    """
    Gauss-Newton implementation for least-squares fitting of non-linear functions
    
    :param p: initial value(s)
    :param function: user-provided function which takes p (and additional arguments) as input
    :param jac_fun: user-provided Jacobian function which takes p (and additional arguments) as input
    :param args: optional arguments passed to function
    :param ftol: relative change in cost function as stop condition
    :param ptol: relative change in independant variables as stop condition
    :param gtol: maximum gradient tolerance as stop condition
    :param l: step size damping parameter
    :param max_iter: maximum number of iterations
    :return: list of results
    """

    if len(args) > 0:
        # pass optional arguments to function
        fun = lambda p: function(p, *args)
    else:
        fun = function

    if jac_function is None:
        # use numerical Jacobian if analytical is not provided
        jac_fun = lambda p: jacobian_approx_t(p, f=fun)
    else:
        jac_fun = lambda p: jac_function(p, *args)

    f = fun(p)
    j = jac_fun(p)
    g = torch.matmul(j.T, f)
    H = torch.matmul(j.T, j)
    p_list = [p]
    while len(p_list) < max_iter:
        h = -l*torch.linalg.lstsq(H, g, rcond=None, driver=None)[0]
        p = p + h
        p_list.append(p.clone())
        f_prev = f.clone()
        f = fun(p)
        j = jac_fun(p)
        g = torch.matmul(j.T, f)
        H = torch.matmul(j.T, j)

        # stop conditions
        gcon = max(abs(g)) < gtol
        pcon = (h**2).sum()**.5 < ptol*(ptol + (p**2).sum()**.5)
        fcon = ((f_prev-f)**2).sum() < ((ftol*f)**2).sum()
        if gcon or pcon or fcon:
            break

    return p_list
//  [python]
def __find_concentration(self, OD, HE):
        # rows correspond to channels (RGB), columns to OD values
        Y = OD.T

        # determine concentrations of the individual stains
        if not self.updated_lstsq:
            return torch.lstsq(Y, HE)[0][:2]
    
        return torch.lstsq(HE, Y)[0]
//  [python]
def __find_concentration(self, OD, HE):
        # rows correspond to channels (RGB), columns to OD values
        Y = OD.T

        # determine concentrations of the individual stains
        if not self.updated_lstsq:
            return torch.lstsq(Y, HE)[0][:2]
    
        return torch.linalg.lstsq(HE, Y)[0]
//  [python]
def alternating_projections_step(self):
        m = self.y.size(-1)

        def f(z):
            z_comp = z[:m] + z[m:] * 1j
            phased = phase(self.A @ torch.solve(z_comp, self.A.T @ self.A)[0])
            return torch.cat([self.y * phased.real, self.y * phased.imag])

        return f
//  [python]
def alternating_projections_step(self):
        m = self.y.size(-1)

        def f(z):
            z_comp = z[:m] + z[m:] * 1j
            phased = phase(self.A @ torch.linalg.lstsq(self.A, z_comp).solution)
            return torch.cat([self.y * phased.real, self.y * phased.imag])

        return f
//  [python]
def __call__(self, xs, inds=None):
        xs = xs.cpu()
        if inds is None:
            inds = range(xs.shape[1])
        else:
            if max(inds) >= xs.shape[1] or min(inds) < 0:
                raise ValueError("inds contain indices that are not defined")

        preds = []
        for i in inds:
            if i <= self.H:
                preds.append(torch.zeros_like(xs[:, 0]))  # predict zero for first point
                continue
            train_zs = torch.zeros(xs.shape[0], i-self.H, int(xs.shape[2]*self.H))
            for j in range(i-self.H):
                train_zs[:,j,:] = xs[:,j:j+self.H,:].reshape(xs.shape[0],-1)
            train_ys = xs[:, self.H:i]
            test_z = xs[:, i - self.H + 1 : i + 1].reshape(xs.shape[0],1,-1)

            ws, _, _, _ = torch.solve(train_ys.t() @ train_zs, train_zs.t() @ train_zs)
            ws = ws.transpose(0,1)

            pred = test_z @ ws

            preds.append(pred[:, 0, :])

        return torch.stack(preds, dim=1)
//  [python]
def __call__(self, xs, inds=None):
        xs = xs.cpu()
        if inds is None:
            inds = range(xs.shape[1])
        else:
            if max(inds) >= xs.shape[1] or min(inds) < 0:
                raise ValueError("inds contain indices that are not defined")

        preds = []
        for i in inds:
            if i <= self.H:
                preds.append(torch.zeros_like(xs[:, 0]))  # predict zero for first point
                continue
            train_zs = torch.zeros(xs.shape[0], i-self.H, int(xs.shape[2]*self.H))
            for j in range(i-self.H):
                train_zs[:,j,:] = xs[:,j:j+self.H,:].reshape(xs.shape[0],-1)
            train_ys = xs[:, self.H:i]
            test_z = xs[:, i - self.H + 1 : i + 1].reshape(xs.shape[0],1,-1)

            ws, _, _, _ = torch.linalg.lstsq(
                train_zs, train_ys, driver=self.driver
            )

            pred = test_z @ ws

            preds.append(pred[:, 0, :])

        return torch.stack(preds, dim=1)
//  [python]
def calculate_self_energy(self, Q, E):
        
        '''
        computes the per-atom contribution from Q towards the property E
        '''
        nmol = len(Q)
        
        natom_counts = torch.zeros(nmol, dtype=torch.int)
        
        for i in range(nmol):
            natom_counts[i] = Q[i].shape[0]
            
        max_atoms = natom_counts.max().item()
        
        paddedZ = torch.zeros(nmol, max_atoms, dtype=torch.int)
        
        X = torch.zeros(nmol, len(self.elements), dtype=torch.float64)
        
        for i in range(nmol):
            paddedZ[i,:natom_counts[i]] = torch.from_numpy(Q[i])

        for i, e in enumerate(self.elements):
            indexes = paddedZ == e
            
            numbers = torch.sum(indexes, dim=1)
            
            X[:, i] = numbers
        
        XTX = torch.matmul(X.T, X)
            
        energies = torch.DoubleTensor(E)
        
        if (self.convert_hartree2kcal()):
            energies *= self.hartree2kcalmol
        
        beta = torch.solve(torch.matmul(X.T, energies[:, None]), XTX)[0][:, 0]
        
        del XTX
        
        species = self.elements.long()
        
        self.self_energy = torch.zeros(torch.max(species) + 1, dtype=torch.float64)
    
        self.self_energy[species] = beta
        
        print ("Self Energies: ", self.self_energy)
//  [python]
def calculate_self_energy(self, Q, E):
        
        '''
        computes the per-atom contribution from Q towards the property E
        '''
        nmol = len(Q)
        
        natom_counts = torch.zeros(nmol, dtype=torch.int)
        
        for i in range(nmol):
            natom_counts[i] = Q[i].shape[0]
            
        max_atoms = natom_counts.max().item()
        
        paddedZ = torch.zeros(nmol, max_atoms, dtype=torch.int)
        
        X = torch.zeros(nmol, len(self.elements), dtype=torch.float64)
        
        for i in range(nmol):
            paddedZ[i,:natom_counts[i]] = torch.from_numpy(Q[i])

        for i, e in enumerate(self.elements):
            indexes = paddedZ == e
            
            numbers = torch.sum(indexes, dim=1)
            
            X[:, i] = numbers
        
        XTX = torch.matmul(X.T, X)
            
        energies = torch.DoubleTensor(E)
        
        if (self.convert_hartree2kcal()):
            energies *= self.hartree2kcalmol
        
        beta = torch.linalg.lstsq(XTX, torch.matmul(X.T, energies[:, None])).solution[:, 0]
        
        del XTX
        
        species = self.elements.long()
        
        self.self_energy = torch.zeros(torch.max(species) + 1, dtype=torch.float64)
    
        self.self_energy[species] = beta
        
        print ("Self Energies: ", self.self_energy)
//  [python]
def _polyfit(potential, quantity, polyorder: int):
    x = torch.stack([potential ** i for i in range(polyorder + 1)], dim=-1) 

    coeffs = torch.lstsq(quantity, x).solution
    return coeffs
//  [python]
def _polyfit(potential, quantity, polyorder: int):
    x = torch.stack([potential ** i for i in range(polyorder + 1)], dim=-1) 

    coeffs = torch.linalg.lstsq(x, quantity).solution
    return coeffs
//  [python]
def lstsq(a: ArrayLike, b: ArrayLike, rcond=None):
    a, b = _atleast_float_2(a, b)
    # NumPy is using gelsd: https://github.com/numpy/numpy/blob/v1.24.0/numpy/linalg/umath_linalg.cpp#L3991
    # on CUDA, only `gels` is available though, so use it instead
    driver = "gels" if a.is_cuda or b.is_cuda else "gelsd"
    return torch.linalg.lstsq(a, b, rcond=rcond, driver=driver)
//  [python]
def solve_lsr(PE_can_list, PE_obs_list):
    rot_mtx_forward = []
    rot_mtx_back = []
    for xyz_can, xyz_obs in zip(PE_can_list, PE_obs_list):
        xyz_can = xyz_can.float()
        xyz_obs = xyz_obs.float()
        # print(xyz_can.shape, xyz_obs.shape) # N, 3 
        rot_mtx_forward.append(torch.solve(torch.matmul(xyz_can.T, xyz_obs), torch.matmul(xyz_can.T, xyz_can)).solution) 
        rot_mtx_back.append(torch.solve(torch.matmul(xyz_obs.T, xyz_can), torch.matmul(xyz_obs.T, xyz_obs)).solution) # 3, 3 
        # print(rot_mtx_forward.shape, rot_mtx_back.shape) 
    rot_mtx_forward = torch.stack(rot_mtx_forward) 
    rot_mtx_back = torch.stack(rot_mtx_back)
    
    return rot_mtx_forward, rot_mtx_back
//  [python]
def solve_lsr(PE_can_list, PE_obs_list):
    rot_mtx_forward = []
    rot_mtx_back = []
    for xyz_can, xyz_obs in zip(PE_can_list, PE_obs_list):
        xyz_can = xyz_can.float()
        xyz_obs = xyz_obs.float()
        # print(xyz_can.shape, xyz_obs.shape) # N, 3 
        rot_mtx_forward.append(torch.linalg.lstsq(xyz_can, xyz_obs).solution) 
        rot_mtx_back.append(torch.linalg.lstsq(xyz_obs, xyz_can).solution) # 3, 3 
        # print(rot_mtx_forward.shape, rot_mtx_back.shape) 
    rot_mtx_forward = torch.stack(rot_mtx_forward) 
    rot_mtx_back = torch.stack(rot_mtx_back)
    
    return rot_mtx_forward, rot_mtx_back
//  [python]
def project_to_null_space(params: torch.Tensor, null_space: torch.Tensor) -> torch.Tensor:
    """
    Project o_s to null_space

    flattened_o_s: (n-params) or (n-models, n-params)
    null_space: (n-nullspace-dims, n-params)

    returns: (1 or n-models, n-params)
    """
    if params.dim() == 1:
        params = params.unsqueeze(0)
    if null_space.shape[1] == params.shape[1] + 1:
        params = F.pad(params, (0, 1), "constant", 1)
    elif null_space.shape[1] != params.shape[1]:
        raise ValueError(f"Shape mismatch {null_space.shape[1]} != {params.shape[1]}")

    return torch.solve(params.unsqueeze(-1), null_space.T.unsqueeze(0))[0].squeeze(-1)
//  [python]
def project_to_null_space(params: torch.Tensor, null_space: torch.Tensor) -> torch.Tensor:
    """
    Project o_s to null_space

    flattened_o_s: (n-params) or (n-models, n-params)
    null_space: (n-nullspace-dims, n-params)

    returns: (1 or n-models, n-params)
    """
    if params.dim() == 1:
        params = params.unsqueeze(0)
    if null_space.shape[1] == params.shape[1] + 1:
        params = F.pad(params, (0, 1), "constant", 1)
    elif null_space.shape[1] != params.shape[1]:
        raise ValueError(f"Shape mismatch {null_space.shape[1]} != {params.shape[1]}")

    return torch.linalg.lstsq(null_space.T.unsqueeze(0), params.unsqueeze(-1)).solution.squeeze(-1)
//  [python]
def lstsq_regression(self, datamodule):
        train_dataset = datamodule.train_dataset
        if self.max_train_points:
            idcs = torch.randperm(len(train_dataset))[: self.max_train_points]
            n_points = len(idcs)
            logger.info(
                f"Calculating weights for regression for relative atomic energy using "
                f"{n_points} randomly sampled points from the training set..."
            )
        else:
            n_points = len(train_dataset)
            idcs = torch.arange(n_points)
            logger.info(
                f"Calculating weights for regression for relative atomic energy using "
                f"all {n_points} points from the training set."
            )
        # read regression variables x (concentration of atoms in molecules +1 column
        # for bias) and regression targets y (energy per atom of molecules)
        x = torch.ones((n_points, len(self.atom_types) + 1))
        y = torch.empty((n_points, 1))
        idx = 0
        for inputs in train_dataset.iter_properties(
            load_properties=[self.target_energy_name],
            load_structure=True,
            indices=idcs.tolist(),
        ):
            x[idx, :-1] = self.get_concentration(inputs)
            y[idx] = inputs[self.target_energy_name] / inputs[properties.n_atoms]
            idx += 1
        # compute LSTSQ solution to obtain weights and bias
        weights = torch.lstsq(y, x).solution
        regression_weights = weights[:-1, 0]
        regression_bias = weights[-1, 0]

        return regression_weights, regression_bias
//  [python]
def lstsq_regression(self, datamodule):
        train_dataset = datamodule.train_dataset
        if self.max_train_points:
            idcs = torch.randperm(len(train_dataset))[: self.max_train_points]
            n_points = len(idcs)
            logger.info(
                f"Calculating weights for regression for relative atomic energy using "
                f"{n_points} randomly sampled points from the training set..."
            )
        else:
            n_points = len(train_dataset)
            idcs = torch.arange(n_points)
            logger.info(
                f"Calculating weights for regression for relative atomic energy using "
                f"all {n_points} points from the training set."
            )
        # read regression variables x (concentration of atoms in molecules +1 column
        # for bias) and regression targets y (energy per atom of molecules)
        x = torch.ones((n_points, len(self.atom_types) + 1))
        y = torch.empty((n_points, 1))
        idx = 0
        for inputs in train_dataset.iter_properties(
            load_properties=[self.target_energy_name],
            load_structure=True,
            indices=idcs.tolist(),
        ):
            x[idx, :-1] = self.get_concentration(inputs)
            y[idx] = inputs[self.target_energy_name] / inputs[properties.n_atoms]
            idx += 1
        # compute LSTSQ solution to obtain weights and bias
        weights = torch.linalg.lstsq(x, y, rcond=1e-8, driver="gelsd").solution
        regression_weights = weights[:-1, 0]
        regression_bias = weights[-1, 0]

        return regression_weights, regression_bias
//  [python]
def anderson_qr_factorization(X, relaxation=1.0, regularization = 0.0):
    # Anderson Acceleration
    # Take a matrix X of iterates such that X[:,i] = g(X[:,i-1])
    # Return acceleration for X[:,-1]

    assert X.ndim==2, "X must be a matrix"
    assert regularization >= 0.0, "regularization for least-squares must be >=0.0"

    # Compute residuals
    DX =  X[:,1:] -  X[:,:-1] # DX[:,i] =  X[:,i+1] -  X[:,i]
    DR = DX[:,1:] - DX[:,:-1] # DR[:,i] = DX[:,i+1] - DX[:,i] = X[:,i+2] - 2*X[:,i+1] + X[:,i]

    if regularization == 0.0:
       # solve unconstrained least-squares problem
       gamma = torch.solve(DR.t() @ DX[:, -1], DR.t() @ DR).solution
    else:
       # solve augmented least-squares for Tykhonov regularization
       rhs = DX[:,-1]
       expanded_rhs    = torch.cat( (rhs, torch.zeros(DR.size(1))) )
       expanded_matrix = torch.cat( (DR, torch.sqrt(torch.tensor(regularization)) * torch.eye(DR.size(1))) )
       gamma = torch.solve(expanded_matrix.t() @ expanded_rhs, expanded_matrix.t() @ expanded_matrix).solution

    # compute acceleration
    extr = X[:,-2] + DX[:,-1] - (DX[:,:-1]+DR)@gamma

    if relaxation!=1:
        assert relaxation>0, "relaxation must be positive"
        # compute solution of the contraint optimization problem s.t. gamma = X[:,1:]@alpha
        alpha = torch.zeros(gamma.numel()+1).to(DX.device)
        alpha[0]    = gamma[0]
        alpha[1:-1] = gamma[1:] - gamma[:-1]
        alpha[-1]   = 1 - gamma[-1]
        extr = relaxation*extr + (1-relaxation)*X[:,:-1]@alpha

    return extr
//  [python]
def anderson_qr_factorization(X, relaxation=1.0, regularization = 0.0):
    # Anderson Acceleration
    # Take a matrix X of iterates such that X[:,i] = g(X[:,i-1])
    # Return acceleration for X[:,-1]

    assert X.ndim==2, "X must be a matrix"
    assert regularization >= 0.0, "regularization for least-squares must be >=0.0"

    # Compute residuals
    DX =  X[:,1:] -  X[:,:-1] # DX[:,i] =  X[:,i+1] -  X[:,i]
    DR = DX[:,1:] - DX[:,:-1] # DR[:,i] = DX[:,i+1] - DX[:,i] = X[:,i+2] - 2*X[:,i+1] + X[:,i]

    if regularization == 0.0:
       # solve unconstrained least-squares problem
       gamma = torch.linalg.lstsq(DR, DX[:, -1]).solution
    else:
       # solve augmented least-squares for Tykhonov regularization
       rhs = DX[:,-1]
       expanded_rhs    = torch.cat( (rhs, torch.zeros(DR.size(1))) )
       expanded_matrix = torch.cat( (DR, torch.sqrt(torch.tensor(regularization)) * torch.eye(DR.size(1))) )
       gamma = torch.linalg.lstsq(expanded_matrix, expanded_rhs).solution

    # compute acceleration
    extr = X[:,-2] + DX[:,-1] - (DX[:,:-1]+DR)@gamma

    if relaxation!=1:
        assert relaxation>0, "relaxation must be positive"
        # compute solution of the contraint optimization problem s.t. gamma = X[:,1:]@alpha
        alpha = torch.zeros(gamma.numel()+1).to(DX.device)
        alpha[0]    = gamma[0]
        alpha[1:-1] = gamma[1:] - gamma[:-1]
        alpha[-1]   = 1 - gamma[-1]
        extr = relaxation*extr + (1-relaxation)*X[:,:-1]@alpha

    return extr
//  [python]
def compute_average_E0(
    dataset,
    symbols: Optional[List[str]]=None,
):
    """
    Function to compute the average interaction energy of each chemical element
    returns dictionary of E0s
    """
    numbers = [atomic_numbers[s] for s in symbols]
    len_train = len(dataset)
    num_elements = len(numbers)
    
    A = torch.zeros((len_train, num_elements))
    B = torch.zeros((len_train,))
    
    for i in range(len_train):
        B[i] = dataset[i][properties.energy]
        for j, z in enumerate(numbers):
            A[i, j] = torch.count_nonzero(dataset[i][properties.Z] == z)
    atomic_energies_dict = {z: 0.0 for z in numbers}
    try:
        E0s, _, _, _ = torch.lstsq(B, A)
        E0s = E0s.T
        for i, z in enumerate(numbers):
            atomic_energies_dict[z] = E0s[i].item()
            
    except RuntimeError:
        print(
            "Failed to compute E0s using least squares regression, using the same for all atoms"
        )
    return atomic_energies_dict
//  [python]
def compute_average_E0(
    dataset,
    symbols: Optional[List[str]]=None,
):
    """
    Function to compute the average interaction energy of each chemical element
    returns dictionary of E0s
    """
    numbers = [atomic_numbers[s] for s in symbols]
    len_train = len(dataset)
    num_elements = len(numbers)
    
    A = torch.zeros((len_train, num_elements))
    B = torch.zeros((len_train,))
    
    for i in range(len_train):
        B[i] = dataset[i][properties.energy]
        for j, z in enumerate(numbers):
            A[i, j] = torch.count_nonzero(dataset[i][properties.Z] == z)
    atomic_energies_dict = {z: 0.0 for z in numbers}
    try:
        E0s = torch.linalg.lstsq(A, B, rcond=None)[0]
        for i, z in enumerate(numbers):
            atomic_energies_dict[z] = E0s[i].item()
            
    except torch.linalg.LinAlgError:
        print(
            "Failed to compute E0s using least squares regression, using the same for all atoms"
        )
    return atomic_energies_dict
//  [python]
def infer_meta_from_obs_and_rewards(self, obs: torch.Tensor, reward: torch.Tensor) -> MetaDict:
        print('max reward: ', reward.max().cpu().item())
        print('99 percentile: ', torch.quantile(reward, 0.99).cpu().item())
        print('median reward: ', reward.median().cpu().item())
        print('min reward: ', reward.min().cpu().item())
        print('mean reward: ', reward.mean().cpu().item())
        print('num reward: ', reward.shape[0])

        with torch.no_grad():
            phi = self.feature_learner.feature_net(obs)
        z = torch.solve(torch.matmul(phi.t(), reward), torch.matmul(phi.t(), phi)).solution  # z_dim x 1
        z = math.sqrt(self.cfg.z_dim) * F.normalize(z, dim=0)  # be careful to the dimension
        meta = OrderedDict()
        meta['z'] = z.squeeze().cpu().numpy()
        # self.solved_meta = meta
        return meta
//  [python]
def infer_meta_from_obs_and_rewards(self, obs: torch.Tensor, reward: torch.Tensor) -> MetaDict:
        print('max reward: ', reward.max().cpu().item())
        print('99 percentile: ', torch.quantile(reward, 0.99).cpu().item())
        print('median reward: ', reward.median().cpu().item())
        print('min reward: ', reward.min().cpu().item())
        print('mean reward: ', reward.mean().cpu().item())
        print('num reward: ', reward.shape[0])

        with torch.no_grad():
            phi = self.feature_learner.feature_net(obs)
        z = torch.linalg.lstsq(phi, reward).solution  # z_dim x 1
        z = math.sqrt(self.cfg.z_dim) * F.normalize(z, dim=0)  # be careful to the dimension
        meta = OrderedDict()
        meta['z'] = z.squeeze().cpu().numpy()
        # self.solved_meta = meta
        return meta
//  [python]
def solve_linear_regression(self, X, Y):
        X = X.to(torch.float64)
        Y = Y.to(torch.float64)
        coef = torch.linalg.pinv(X.T @ X) @ X.T @ Y
        # coef = torch.linalg.lstsq(X, Y).solution
        return coef
//  [python]
def solve_linear_regression(self, X, Y):
        X = X.to(torch.float64)
        Y = Y.to(torch.float64)
        # coef = torch.linalg.pinv(X.T @ X) @ X.T @ Y
        coef = torch.linalg.lstsq(X, Y).solution
        return coef
//  [python]
def fit_circle_2d(x, y, w=[]):
    
    A = torch.stack([x, y, torch.ones(len(x))]).T
    b = x**2 + y**2
    
    # Modify A,b for weighted least squares
    if len(w) == len(x):
        W = torch.diag(w)
        A = torch.mm(W,A)
        b = torch.mv(W,b)
    
    # Solve by method of least squares
    c = torch.lstsq(b,A)[0]
    
    # Get circle parameters from solution c
    xc = c[0]/2
    yc = c[1]/2
    r = torch.sqrt(c[2] + xc**2 + yc**2)
    return xc, yc, r
//  [python]
def fit_circle_2d(x, y, w=[]):
    
    A = torch.stack([x, y, torch.ones(len(x))]).T
    b = x**2 + y**2
    
    # Modify A,b for weighted least squares
    if len(w) == len(x):
        W = torch.diag(w)
        A = torch.dot(W,A)
        b = torch.dot(W,b)
    
    # Solve by method of least squares
    c = torch.linalg.lstsq(A,b,rcond=None)[0]
    
    # Get circle parameters from solution c
    xc = c[0]/2
    yc = c[1]/2
    r = torch.sqrt(c[2] + xc**2 + yc**2)
    return xc, yc, r
//  [python]
def controllable_speakers(self, intermediate, z):
        pca = torch.pca_lowrank(intermediate)
        mu = intermediate.mean()
        X = torch.matmul((intermediate - mu), pca[2])
        U = torch.lstsq(X, z)
        return U
//  [python]
def controllable_speakers(self, intermediate, z):
        pca = torch.pca_lowrank(intermediate)
        mu = intermediate.mean()
        X = torch.matmul((intermediate - mu), pca[2])
        U = torch.linalg.lstsq(X, z)
        return U
//  [python]
def matrix_solve_least_squares(self, matrix: TensorType, rhs: TensorType) -> Tuple[TensorType, TensorType, TensorType, TensorType]:
        assert version.parse(torch.__version__) >= version.parse('1.9.0'), "least squares requires PyTorch >= 1.9.0"
        matrix, rhs = self.auto_cast(matrix, rhs)
        solution, residuals, rank, singular_values = torch.linalg.solve(matrix.t() @ matrix, matrix.t() @ rhs)
        return solution, residuals, rank, singular_values
//  [python]
def matrix_solve_least_squares(self, matrix: TensorType, rhs: TensorType) -> Tuple[TensorType, TensorType, TensorType, TensorType]:
        assert version.parse(torch.__version__) >= version.parse('1.9.0'), "least squares requires PyTorch >= 1.9.0"
        matrix, rhs = self.auto_cast(matrix, rhs)
        solution, residuals, rank, singular_values = torch.linalg.lstsq(matrix, rhs)
        return solution, residuals, rank, singular_values
//  [python]
def intersect_skew_lines_high_dim(p, r, mask=None):
    # Implements https://en.wikipedia.org/wiki/Skew_lines In more than two dimensions
    dim = p.shape[-1]
    # make sure the heading vectors are l2-normed
    if mask is None:
        mask = torch.ones_like(p[..., 0])
    r = torch.nn.functional.normalize(r, dim=-1)

    eye = torch.eye(dim, device=p.device, dtype=p.dtype)[None, None]
    I_min_cov = (eye - (r[..., None] * r[..., None, :])) * mask[..., None, None]
    sum_proj = I_min_cov.matmul(p[..., None]).sum(dim=-3)
    p_intersect = torch.solve(sum_proj, I_min_cov.sum(dim=-3)).solution[..., 0]

    if torch.any(torch.isnan(p_intersect)):
        print(p_intersect)
        raise ValueError(f"p_intersect is NaN")

    return p_intersect, r
//  [python]
def intersect_skew_lines_high_dim(p, r, mask=None):
    # Implements https://en.wikipedia.org/wiki/Skew_lines In more than two dimensions
    dim = p.shape[-1]
    # make sure the heading vectors are l2-normed
    if mask is None:
        mask = torch.ones_like(p[..., 0])
    r = torch.nn.functional.normalize(r, dim=-1)

    eye = torch.eye(dim, device=p.device, dtype=p.dtype)[None, None]
    I_min_cov = (eye - (r[..., None] * r[..., None, :])) * mask[..., None, None]
    sum_proj = I_min_cov.matmul(p[..., None]).sum(dim=-3)
    p_intersect = torch.linalg.lstsq(I_min_cov.sum(dim=-3), sum_proj).solution[..., 0]

    if torch.any(torch.isnan(p_intersect)):
        print(p_intersect)
        raise ValueError(f"p_intersect is NaN")

    return p_intersect, r
//  [python]
def add_batch_update_u(self, inp, out):
        inps = inp.view(inp.shape[0] * inp.shape[1], inp.shape[2])
        outs = out.view(out.shape[0] * out.shape[1], out.shape[2])
        new_w = torch.matmul(self.truc_u, torch.matmul(self.truc_sigma, self.truc_v))
        new_output = inps.matmul(new_w.t())
        self.error = torch.sqrt(torch.sum((outs - new_output)**2)).item() / torch.norm(outs, p='fro').item()
        # print(f"truncted error: {self.error}")
        x =  torch.matmul(torch.matmul(inps, self.truc_v.T), self.truc_sigma)
        self.updated_uT = torch.solve(torch.matmul(x.T, outs), torch.matmul(x.T, x)).solution
        updated_output = torch.matmul(torch.matmul(torch.matmul(inps, self.truc_v.T), self.truc_sigma), self.updated_uT)
        self.updated_error = torch.sqrt(torch.sum((outs - updated_output)**2)).item() / torch.norm(outs, p='fro').item()
        # print(f"updated error: {self.updated_error}")
        inps = outs = new_output = updated_output = x = new_w = None
        del inps, outs, new_output, updated_output, x, new_w
        torch.cuda.empty_cache()
//  [python]
def add_batch_update_u(self, inp, out):
        inps = inp.view(inp.shape[0] * inp.shape[1], inp.shape[2])
        outs = out.view(out.shape[0] * out.shape[1], out.shape[2])
        new_w = torch.matmul(self.truc_u, torch.matmul(self.truc_sigma, self.truc_v))
        new_output = inps.matmul(new_w.t())
        self.error = torch.sqrt(torch.sum((outs - new_output)**2)).item() / torch.norm(outs, p='fro').item()
        # print(f"truncted error: {self.error}")
        x =  torch.matmul(torch.matmul(inps, self.truc_v.T), self.truc_sigma)
        self.updated_uT = torch.linalg.lstsq(x,outs).solution
        updated_output = torch.matmul(torch.matmul(torch.matmul(inps, self.truc_v.T), self.truc_sigma), self.updated_uT)
        self.updated_error = torch.sqrt(torch.sum((outs - updated_output)**2)).item() / torch.norm(outs, p='fro').item()
        # print(f"updated error: {self.updated_error}")
        inps = outs = new_output = updated_output = x = new_w = None
        del inps, outs, new_output, updated_output, x, new_w
        torch.cuda.empty_cache()