jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def match_channel_colors(im_ref, im_q, im_test):

    im_ref_reshape = im_ref.view(*im_ref.shape[:2], -1)
    im_q_reshape = im_q.view(*im_q.shape[:2], -1)
    im_test_reshape = im_test.view(*im_test.shape[:2], -1)
    # Estimate color transformation matrix by minimizing the least squares error

    im_t_conv_list = []
    for i in range(im_ref.size(1)):
        c_mat_all = []
        for ir_batch, iq_batch in zip(im_ref_reshape[:, i:i+1, :], im_q_reshape[:, i:i+1, :]):
            c = torch.linalg.lstsq(ir_batch.t(), iq_batch.t())
            c = c.solution[:1]
            c_mat_all.append(c)

        c_mat = torch.stack(c_mat_all, dim=0)
        # Apply the transformation to test image
        im_t_conv = torch.matmul(im_test_reshape[:, i:i+1, :].permute(0, 2, 1), c_mat).permute(0, 2, 1)
        im_t_conv = im_t_conv.view(*im_t_conv.shape[:2], *im_test.shape[-2:])
        im_t_conv_list.append(im_t_conv)

    im_t_conv = torch.cat(im_t_conv_list, dim=1)

    return im_t_conv
//  [python]
def _compute_coefficients(self):
        """
        Compute the coefficients of the polynomial spline, and set them into the internal linear layer for storing.
        """
        # Treat each action dimension separately
        for idx_act in range(self.env_spec.act_space.flat_dim):
            # Get the feature matrices for both points in time
            feats = to.cat([self._compute_feats(self._t_init), self._compute_feats(self._t_end)], dim=0)

            # Solve least squares problem
            coeffs = to.lstsq(self.conds[:, idx_act], feats).solution

            # Store
            self.coeffs[:, idx_act] = coeffs.squeeze()
//  [python]
def _compute_coefficients(self):
        """
        Compute the coefficients of the polynomial spline, and set them into the internal linear layer for storing.
        """
        # Treat each action dimension separately
        for idx_act in range(self.env_spec.act_space.flat_dim):
            # Get the feature matrices for both points in time
            feats = to.cat([self._compute_feats(self._t_init), self._compute_feats(self._t_end)], dim=0)

            # Solve least squares problem
            coeffs = to.linalg.lstsq(self.conds[:, idx_act], feats).solution

            # Store
            self.coeffs[:, idx_act] = coeffs.squeeze()
//  [python]
def compute_coefficients(self):
        """
        Compute the coefficients of the polynomial spline, and set them into the internal linear layer for storing.
        """
        # Treat each action dimension separately
        for idx_act in range(self.act_space_flat_dim):
            # Get the feature matrices for both points in time
            feats = to.cat([self.compute_feats(self.t_init), self.compute_feats(self.t_end)], dim=0)

            # Solve least squares problem
            coeffs = to.lstsq(self.conds[:, idx_act], feats).solution

            # Store
            self.coeffs[:, idx_act] = coeffs.squeeze()
//  [python]
def compute_coefficients(self):
        """
        Compute the coefficients of the polynomial spline, and set them into the internal linear layer for storing.
        """
        # Treat each action dimension separately
        for idx_act in range(self.act_space_flat_dim):
            # Get the feature matrices for both points in time
            feats = to.cat([self.compute_feats(self.t_init), self.compute_feats(self.t_end)], dim=0)

            # Solve least squares problem
            coeffs = to.linalg.lstsq(self.conds[:, idx_act], feats).solution

            # Store
            self.coeffs[:, idx_act] = coeffs.squeeze()
//  [python]
def ProjectFunctionNeRFSparse(
    order: int,
    spherical_func: Callable,
    sample_count: int,
    device="cpu",
):
    assert order >= 0, "Order must be at least zero."
    assert sample_count > 0, "Sample count must be at least one."
    C = 3  # rgb channels

    # generate sample_count uniformly and stratified samples over the sphere
    # See http://www.bogotobogo.com/Algorithms/uniform_distribution_sphere.php
    theta, phi = spherical_uniform_sampling(sample_count, device=device)
    dirs = spher2cart(theta, phi)  # [sample_count, 3]

    # evaluate the analytic function for the current spherical coords
    func_value, others = spherical_func(dirs)  # func_value [batch_size, sample_count, C]

    batch_size = func_value.shape[0]

    coeff_count = GetCoefficientCount(order)
    basis_vals = torch.empty(
        [sample_count, coeff_count], dtype=torch.float32
    ).to(device)

    # evaluate the SH basis functions up to band O, scale them by the
    # function's value and accumulate them over all generated samples
    for l in range(order + 1):  # end inclusive
        for m in range(-l, l + 1):  # end inclusive
            basis_vals[:, GetIndex(l, m)] = EvalSH(l, m, dirs)

    basis_vals = basis_vals.view(
           sample_count, coeff_count) # [sample_count, coeff_count]
    func_value = func_value.transpose(0, 1).reshape(
           sample_count, batch_size * C) # [sample_count, batch_size * C]
    soln = torch.lstsq(func_value, basis_vals).solution[:basis_vals.size(1)]
    soln = soln.T.reshape(batch_size, C, -1)
    return soln, others
//  [python]
def ProjectFunctionNeRFSparse(
    order: int,
    spherical_func: Callable,
    sample_count: int,
    device="cpu",
):
    assert order >= 0, "Order must be at least zero."
    assert sample_count > 0, "Sample count must be at least one."
    C = 3  # rgb channels

    # generate sample_count uniformly and stratified samples over the sphere
    # See http://www.bogotobogo.com/Algorithms/uniform_distribution_sphere.php
    theta, phi = spherical_uniform_sampling(sample_count, device=device)
    dirs = spher2cart(theta, phi)  # [sample_count, 3]

    # evaluate the analytic function for the current spherical coords
    func_value, others = spherical_func(dirs)  # func_value [batch_size, sample_count, C]

    batch_size = func_value.shape[0]

    coeff_count = GetCoefficientCount(order)
    basis_vals = torch.empty(
        [sample_count, coeff_count], dtype=torch.float32
    ).to(device)

    # evaluate the SH basis functions up to band O, scale them by the
    # function's value and accumulate them over all generated samples
    for l in range(order + 1):  # end inclusive
        for m in range(-l, l + 1):  # end inclusive
            basis_vals[:, GetIndex(l, m)] = EvalSH(l, m, dirs)

    basis_vals = basis_vals.view(
           sample_count, coeff_count) # [sample_count, coeff_count]
    func_value = func_value.transpose(0, 1).reshape(
           sample_count, batch_size * C) # [sample_count, batch_size * C]
    soln = torch.linalg.lstsq(func_value, basis_vals).solution[:basis_vals.size(1)]
    soln = soln.T.reshape(batch_size, C, -1)
    return soln, others
//  [python]
def _make_coeffs(window_length, order):
        half_length, rem = divmod(window_length, 2)
        if rem == 0:
            raise ValueError("window_length must be odd.")

        idx = torch.arange(window_length - half_length - 1, -half_length - 1, -1, dtype=torch.float, device="cpu")
        a = idx ** torch.arange(order + 1, dtype=torch.float, device="cpu").reshape(-1, 1)
        y = torch.zeros(order + 1, dtype=torch.float, device="cpu")
        y[0] = 1.0
        return (
            torch.lstsq(y, a).solution.squeeze()  # type: ignore
            if not pytorch_after(1, 11)
            else torch.linalg.lstsq(a, y).solution.squeeze()
        )
//  [python]
def _make_coeffs(window_length, order):
        half_length, rem = divmod(window_length, 2)
        if rem == 0:
            raise ValueError("window_length must be odd.")

        idx = torch.arange(window_length - half_length - 1, -half_length - 1, -1, dtype=torch.float, device="cpu")
        a = idx ** torch.arange(order + 1, dtype=torch.float, device="cpu").reshape(-1, 1)
        y = torch.zeros(order + 1, dtype=torch.float, device="cpu")
        y[0] = 1.0
        return torch.linalg.lstsq(a, y).solution.squeeze()
//  [python]
def scale_compute(gt, prediction):
    scale, _ = torch.lstsq(gt.flatten().unsqueeze(1), prediction.flatten().unsqueeze(1))
    return scale[0, 0].clone().detach()
//  [python]
def scale_compute(gt, prediction):
    scale, _ = torch.linalg.lstsq(gt.flatten().unsqueeze(1), prediction.flatten().unsqueeze(1))
    return scale[0, 0].clone().detach()
//  [python]
def test_lstsq_removed_error(self, device):
        a = make_tensor(5, 5, device=device, dtype=torch.float32)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            torch.lstsq(a, a)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            a.lstsq(a)
//  [python]
def test_lstsq_removed_error(self, device):
        a = make_tensor(5, 5, device=device, dtype=torch.float32)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            torch.linalg.lstsq(a, a)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            a.lstsq(a)
//  [python]
def _get_nu(self, fY, hY):
        """Compute nu (ie lambda) if not provided by the problem's solver.
        That is, solve: hY^T nu = fY^T.
        """
        p = hY.size(1)
        nu = fY.new_zeros(self.b, p)
        for i in range(self.b):  # loop over batch
            solution, _ = torch.lstsq(fY[i, :].unsqueeze(-1), hY[i, :, :].t())      #ToDo: This function has been replaced with torch.linalg.lstsq
            # Htemp = hY[i, :, :].t()
            # ftemp = fY[i, :].unsqueeze(-1)
            # solution = torch.linalg.lstsq(Htemp, ftemp).solution
            nu[i, :] = solution[:p, :].squeeze()  # extract first p values
            # nu[i, :] = solution.squeeze()
        return nu
//  [python]
def _get_nu(self, fY, hY):
        """Compute nu (ie lambda) if not provided by the problem's solver.
        That is, solve: hY^T nu = fY^T.
        """
        p = hY.size(1)
        nu = fY.new_zeros(self.b, p)
        for i in range(self.b):  # loop over batch
            solution, _ = torch.linalg.lstsq(hY[i, :, :].t(), fY[i, :].unsqueeze(-1))
            nu[i, :] = solution[:p, :].squeeze()  # extract first p values
        return nu
//  [python]
def run_per_epoch_evaluations(self, data_loader):
        print('fitting linear probe')
        representations = []
        targets = []
        for batch in data_loader:
            batch = [element.to(self.device) for element in batch]
            loss, view2d, view3d = self.process_batch(batch, optim=None)
            representations.append(view2d)
            targets.append(batch[-1])
            if len(representations) * len(view2d) >= self.args.linear_probing_samples:
                break
        representations = torch.cat(representations, dim=0)
        targets = torch.cat(targets, dim=0)
        if len(representations) >= representations.shape[-1]:
            X, _ = torch.lstsq(targets, representations)
            X, _ = torch.lstsq(targets, representations)
            sol = X[:representations.shape[-1]]
            pred = representations @ sol
            mean_absolute_error = (pred - targets).abs().mean()
            self.writer.add_scalar('linear_probe_mae', mean_absolute_error.item(), self.optim_steps)
        else:
            raise ValueError(
                f'We have less linear_probing_samples {len(representations)} than the metric dimension {representations.shape[-1]}. Linear probing cannot be used.')

        print('finish fitting linear probe')
//  [python]
def run_per_epoch_evaluations(self, data_loader):
        print('fitting linear probe')
        representations = []
        targets = []
        for batch in data_loader:
            batch = [element.to(self.device) for element in batch]
            loss, view2d, view3d = self.process_batch(batch, optim=None)
            representations.append(view2d)
            targets.append(batch[-1])
            if len(representations) * len(view2d) >= self.args.linear_probing_samples:
                break
        representations = torch.cat(representations, dim=0)
        targets = torch.cat(targets, dim=0)
        if len(representations) >= representations.shape[-1]:
            X, _ = torch.linalg.lstsq(targets, representations)
            X, _ = torch.linalg.lstsq(targets, representations)
            sol = X[:representations.shape[-1]]
            pred = representations @ sol
            mean_absolute_error = (pred - targets).abs().mean()
            self.writer.add_scalar('linear_probe_mae', mean_absolute_error.item(), self.optim_steps)
        else:
            raise ValueError(
                f'We have less linear_probing_samples {len(representations)} than the metric dimension {representations.shape[-1]}. Linear probing cannot be used.')

        print('finish fitting linear probe')
//  [python]
def lyapunov_lstsq(A, C):
    """Slow explicit solution to least squares Lyapunov in kronecker expanded form."""
    n, n = A.shape
    ii = torch.eye(n)
    sol = torch.lstsq(u.vec(C), kron(ii, A) + kron(A.t(), ii))[0]
    return unvec(sol, n)
//  [python]
def lyapunov_lstsq(A, C):
    """Slow explicit solution to least squares Lyapunov in kronecker expanded form."""
    n, n = A.shape
    ii = torch.eye(n)
    sol = torch.linalg.lstsq(u.vec(C), kron(ii, A) + kron(A.t(), ii))[0]
    return unvec(sol, n)
//  [python]
def intersect(self, P0, P1):
        """P0 and P1 are NxD arrays defining N lines.
        D is the dimension of the space. This function 
        returns the least squares intersection of the N
        lines from the system given by eq. 13 in 
        http://cal.cs.illinois.edu/~johannes/research/LS_line_intersect.pdf.
        """

        # generate all line direction vectors
        n = (P1 - P0) / torch.norm(P1 - P0, dim=1).unsqueeze(1)  # normalized

        # generate the array of all projectors
        projs = torch.eye(n.shape[1]).float().cuda() - n.unsqueeze(2) * n.unsqueeze(1)  # I - n*n.T

        # generate R matrix and q vector
        R = projs.sum(0)
        q = (projs @ P0.unsqueeze(2)).sum(0)

        # solve the least squares problem for the
        # intersection point p: Rp = q
        p = torch.lstsq(q, R)[0]

        return p
//  [python]
def intersect(self, P0, P1):
        """P0 and P1 are NxD arrays defining N lines.
        D is the dimension of the space. This function 
        returns the least squares intersection of the N
        lines from the system given by eq. 13 in 
        http://cal.cs.illinois.edu/~johannes/research/LS_line_intersect.pdf.
        """

        # generate all line direction vectors
        n = (P1 - P0) / torch.norm(P1 - P0, dim=1).unsqueeze(1)  # normalized

        # generate the array of all projectors
        projs = torch.eye(n.shape[1]).float().cuda() - n.unsqueeze(2) * n.unsqueeze(1)  # I - n*n.T

        # generate R matrix and q vector
        R = projs.sum(0)
        q = (projs @ P0.unsqueeze(2)).sum(0)

        # solve the least squares problem for the
        # intersection point p: Rp = q
        p = torch.linalg.lstsq(R, q)[0]

        return p
//  [python]
def fit_coeff(self, x, weights, y_actual):
        '''
            Use LSE to solve for coeff: y_actual = coeff * (weighted)x
                  x.shape: n_cases * n_in
            weights.shape: n_cases * n_rules
            [ coeff.shape: n_rules * n_out * (n_in+1) ]
                  y.shape: n_cases * n_out
        '''
        # Append 1 to each list of input vals, for the constant term:
        #x_plus = torch.cat([x, torch.ones(x.shape[0], 1)], dim=1)
        x_plus = torch.ones(len(x), 1)

        # Shape of weighted_x is n_cases * n_rules * (n_in+1)
        weighted_x = torch.einsum('bp, bq -> bpq', weights, x_plus)

        # Can't have value 0 for weights, or LSE won't work:
        weighted_x[weighted_x == 0] = 1e-12
        # Squash x and y down to 2D matrices for gels:
        weighted_x_2d = weighted_x.view(weighted_x.shape[0], -1)
        y_actual_2d = y_actual.view(y_actual.shape[0], -1)

        # Use gels to do LSE, then pick out the solution rows:
        try:
            #coeff_2d, _ = torch.gels(y_actual_2d, weighted_x_2d)
            coeff_2d, _ = torch.lstsq(y_actual_2d, weighted_x_2d)
        except RuntimeError as e:
            print('Internal error in gels', e)
            print('Weights are:', weighted_x)
            raise e
        coeff_2d = coeff_2d[0:weighted_x_2d.shape[1]]
        # Reshape to 3D tensor: divide by rules, n_in+1, then swap last 2 dims

        #self.coeff = coeff_2d.view(weights.shape[1], x.shape[1]+1, -1).transpose(1, 2)

        self.coeff = coeff_2d.view(weights.shape[1], self.d_out, -1).transpose(1, 2)
//  [python]
def fit_coeff(self, x, weights, y_actual):
        '''
            Use LSE to solve for coeff: y_actual = coeff * (weighted)x
                  x.shape: n_cases * n_in
            weights.shape: n_cases * n_rules
            [ coeff.shape: n_rules * n_out * (n_in+1) ]
                  y.shape: n_cases * n_out
        '''
        # Append 1 to each list of input vals, for the constant term:
        #x_plus = torch.cat([x, torch.ones(x.shape[0], 1)], dim=1)
        x_plus = torch.ones(len(x), 1)

        # Shape of weighted_x is n_cases * n_rules * (n_in+1)
        weighted_x = torch.einsum('bp, bq -> bpq', weights, x_plus)

        # Can't have value 0 for weights, or LSE won't work:
        weighted_x[weighted_x == 0] = 1e-12
        # Squash x and y down to 2D matrices for gels:
        weighted_x_2d = weighted_x.view(weighted_x.shape[0], -1)
        y_actual_2d = y_actual.view(y_actual.shape[0], -1)

        # Use gels to do LSE, then pick out the solution rows:
        try:
            #coeff_2d, _ = torch.gels(y_actual_2d, weighted_x_2d)
            coeff_2d, _ = torch.linalg.lstsq(y_actual_2d, weighted_x_2d)
        except RuntimeError as e:
            print('Internal error in gels', e)
            print('Weights are:', weighted_x)
            raise e
        coeff_2d = coeff_2d[0:weighted_x_2d.shape[1]]
        # Reshape to 3D tensor: divide by rules, n_in+1, then swap last 2 dims

        #self.coeff = coeff_2d.view(weights.shape[1], x.shape[1]+1, -1).transpose(1, 2)

        self.coeff = coeff_2d.view(weights.shape[1], self.d_out, -1).transpose(1, 2)
//  [python]
def lstsq(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:  # pragma: no cover
    """
    Like torch.linalg.lstsq, tries to return X
    such that AX=B.
    """
    if hasattr(torch, "linalg") and hasattr(torch.linalg, "lstsq"):
        # PyTorch version >= 1.9
        return torch.linalg.lstsq(A, B).solution

    solution = torch.lstsq(B, A).solution
    if A.shape[1] < A.shape[0]:
        return solution[: A.shape[1]]
    return solution
//  [python]
def lstsq(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:  # pragma: no cover
    """
    Like torch.linalg.lstsq, tries to return X
    such that AX=B.
    """
    if hasattr(torch, "linalg") and hasattr(torch.linalg, "lstsq"):
        # PyTorch version >= 1.9
        return torch.linalg.lstsq(A, B).solution

    solution = torch.linalg.lstsq(B, A).solution
    if A.shape[1] < A.shape[0]:
        return solution[: A.shape[1]]
    return solution
//  [python]
def lstsq(input, A):
    solution = []
    QR = []
    for i, a in zip(input, A):
        s, q = torch.lstsq(i, a)
        solution += [s]
        QR += [q]
    solution = torch.stack(solution)
    QR = torch.stack(QR)
    return solution, QR
//  [python]
def lstsq(input, A):
    solution = []
    QR = []
    for i, a in zip(input, A):
        s, q = torch.linalg.lstsq(i, a)
        solution += [s]
        QR += [q]
    solution = torch.stack(solution)
    QR = torch.stack(QR)
    return solution, QR
//  [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.lstsq(b_matrix, a_matrix)[0]

    output: List[float] = res.squeeze(1).tolist()
    return output
//  [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)[0]

    output: List[float] = res.squeeze(1).tolist()
    return output
//  [python]
def fit(self, episodes):
        # sequence_length * batch_size x feature_size
        featmat = self._feature(episodes).view(-1, self.feature_size)
        # sequence_length * batch_size x 1
        returns = episodes.returns.view(-1, 1)

        reg_coeff = self._reg_coeff
        eye = torch.eye(self.feature_size, dtype=torch.float32,
            device=self.linear.weight.device)
        for _ in range(5):
            try:
                # coeffs, _ = torch.gels(
                coeffs, _ = torch.lstsq(
                    torch.matmul(featmat.t(), returns),
                    torch.matmul(featmat.t(), featmat) + reg_coeff * eye
                )
                break
            except RuntimeError:
                reg_coeff += 10
        else:
            raise RuntimeError('Unable to solve the normal equations in '
                '`LinearFeatureBaseline`. The matrix X^T*X (with X the design '
                'matrix) is not full-rank, regardless of the regularization '
                '(maximum regularization: {0}).'.format(reg_coeff))
        self.linear.weight.data = coeffs.data.t()
//  [python]
def fit(self, episodes):
        # sequence_length * batch_size x feature_size
        featmat = self._feature(episodes).view(-1, self.feature_size)
        # sequence_length * batch_size x 1
        returns = episodes.returns.view(-1, 1)

        reg_coeff = self._reg_coeff
        eye = torch.eye(self.feature_size, dtype=torch.float32,
            device=self.linear.weight.device)
        for _ in range(5):
            try:
                coeffs, _ = torch.linalg.lstsq(
                    torch.matmul(featmat.t(), featmat) + reg_coeff * eye,
                    torch.matmul(featmat.t(), returns)
                )
                break
            except RuntimeError:
                reg_coeff += 10
        else:
            raise RuntimeError('Unable to solve the normal equations in '
                '`LinearFeatureBaseline`. The matrix X^T*X (with X the design '
                'matrix) is not full-rank, regardless of the regularization '
                '(maximum regularization: {0}).'.format(reg_coeff))
        self.linear.weight.data = coeffs.data.t()
//  [python]
def find_qbezier_fit(pts, t):
    r"""

    Parameters
    ----------
    pts : torch.Tensor
        of shape [pts_n, spatial_dims_n]
    t : torch.Tensor
        of shape [pts_n]

    Returns
    -------
    control_points : torch.Tensor
        of shape [3, spatial_dims_n]
    error : scalar
        MSE of fit
    """
    T = torch.stack([torch.ones_like(t), t, t.pow(2)]).permute(1, 0)
    del t
    M = T.new_tensor([[1, 0, 0], [-2, 2, 0], [1, -2, 1]])
    M = T.matmul(M)
    del T

    try:
        control_points, _ = torch.lstsq(pts, M)
        del pts, M, _
        error = control_points[3:].pow(2).mean().sqrt().item()
        control_points = control_points[:3]
        return control_points, error
    except RuntimeError:
        return None, np.inf
//  [python]
def find_qbezier_fit(pts, t):
    r"""

    Parameters
    ----------
    pts : torch.Tensor
        of shape [pts_n, spatial_dims_n]
    t : torch.Tensor
        of shape [pts_n]

    Returns
    -------
    control_points : torch.Tensor
        of shape [3, spatial_dims_n]
    error : scalar
        MSE of fit
    """
    T = torch.stack([torch.ones_like(t), t, t.pow(2)]).permute(1, 0)
    del t
    M = T.new_tensor([[1, 0, 0], [-2, 2, 0], [1, -2, 1]])
    M = T.matmul(M)
    del T

    try:
        control_points, _ = torch.linalg.lstsq(pts, M)
        del pts, M, _
        error = control_points[3:].pow(2).mean().sqrt().item()
        control_points = control_points[:3]
        return control_points, error
    except RuntimeError:
        return None, np.inf
//  [python]
def color_shift_remove(fake_image, real_image):
    b, c, h, w = fake_image.shape
    fake_image = fake_image.reshape(c, -1)
    real_image = real_image.reshape(c, -1)
    fake_image_ext = torch.cat([fake_image[[0]] * fake_image[[1]],
                                fake_image[[2]] * fake_image[[1]],
                                fake_image[[2]] * fake_image[[0]],
                                fake_image.pow(2),
                                fake_image,
                                torch.ones_like(fake_image[[0]])
                                ], axis=0)

    x, _ = torch.lstsq(real_image.T, fake_image_ext.T)
    fake_image = torch.mm(fake_image_ext.T, x[:fake_image_ext.shape[0]]).T
    fake_image = fake_image.reshape(b, c, h, w)
    real_image = real_image.reshape(b, c, h, w)

    return fake_image
//  [python]
def color_shift_remove(fake_image, real_image):
    b, c, h, w = fake_image.shape
    fake_image = fake_image.reshape(c, -1)
    real_image = real_image.reshape(c, -1)
    fake_image_ext = torch.cat([fake_image[[0]] * fake_image[[1]],
                                fake_image[[2]] * fake_image[[1]],
                                fake_image[[2]] * fake_image[[0]],
                                fake_image.pow(2),
                                fake_image,
                                torch.ones_like(fake_image[[0]])
                                ], axis=0)

    x, _ = torch.linalg.lstsq(real_image.T, fake_image_ext.T)
    fake_image = torch.mm(fake_image_ext.T, x[:fake_image_ext.shape[0]]).T
    fake_image = fake_image.reshape(b, c, h, w)
    real_image = real_image.reshape(b, c, h, w)

    return fake_image
//  [python]
def __getitem__(self, idx):
        idx1, idx2 = self.combinations[idx]

        shape1 = [
            self.verts_list[idx1],
            self.faces_list[idx1],
            self.frames_list[idx1],
            self.massvec_list[idx1],
            self.L_list[idx1],
            self.evals_list[idx1],
            self.evecs_list[idx1],
            self.gradX_list[idx1],
            self.gradY_list[idx1],
            self.hks_list[idx1],
            self.vts_list[idx1],
            self.names_list[idx1],
        ]

        shape2 = [
            self.verts_list[idx2],
            self.faces_list[idx2],
            self.frames_list[idx2],
            self.massvec_list[idx2],
            self.L_list[idx2],
            self.evals_list[idx2],
            self.evecs_list[idx2],
            self.gradX_list[idx2],
            self.gradY_list[idx2],
            self.hks_list[idx2],
            self.vts_list[idx2],
            self.names_list[idx2],
        ]

        # Compute the ground-truth functional map between the pair
        vts1, vts2 = shape1[10], shape2[10]
        evec_1, evec_2 = shape1[6][:, :self.n_fmap], shape2[6][:, :self.n_fmap]
        evec_1_a, evec_2_a = evec_1[vts1,:], evec_2[vts2,:]
        solve_out = torch.lstsq(evec_2_a, evec_1_a)[0] # TODO replace with torch.linalg version in future torch
        C_gt = solve_out[:evec_1_a.size(-1)].t()
        resids = solve_out[evec_1_a.size(-1):]

        # Alternately, do it with numpy instead:
        # solve_out = np.linalg.lstsq(toNP(evec_1_a), toNP(evec_2_a), rcond=None)
        # C_gt = solve_out[0]
        # C_gt = torch.Tensor(C_gt.T)
        # resids = torch.tensor(solve_out[1])
        
        return (shape1, shape2, C_gt)
//  [python]
def __getitem__(self, idx):
        idx1, idx2 = self.combinations[idx]

        shape1 = [
            self.verts_list[idx1],
            self.faces_list[idx1],
            self.frames_list[idx1],
            self.massvec_list[idx1],
            self.L_list[idx1],
            self.evals_list[idx1],
            self.evecs_list[idx1],
            self.gradX_list[idx1],
            self.gradY_list[idx1],
            self.hks_list[idx1],
            self.vts_list[idx1],
            self.names_list[idx1],
        ]

        shape2 = [
            self.verts_list[idx2],
            self.faces_list[idx2],
            self.frames_list[idx2],
            self.massvec_list[idx2],
            self.L_list[idx2],
            self.evals_list[idx2],
            self.evecs_list[idx2],
            self.gradX_list[idx2],
            self.gradY_list[idx2],
            self.hks_list[idx2],
            self.vts_list[idx2],
            self.names_list[idx2],
        ]

        # Compute the ground-truth functional map between the pair
        vts1, vts2 = shape1[10], shape2[10]
        evec_1, evec_2 = shape1[6][:, :self.n_fmap], shape2[6][:, :self.n_fmap]
        evec_1_a, evec_2_a = evec_1[vts1,:], evec_2[vts2,:]
        solve_out = torch.linalg.lstsq(evec_2_a, evec_1_a)
        C_gt = solve_out[0][:evec_1_a.size(-1)].t()
        resids = solve_out[1]
        
        return (shape1, shape2, C_gt)
//  [python]
def fit_coeff(self, x, weights, y_actual):
        '''
            Use LSE to solve for coeff: y_actual = coeff * (weighted)x
                  x.shape: n_cases * n_in
            weights.shape: n_cases * n_rules
            [ coeff.shape: n_rules * n_out * (n_in+1) ]
                  y.shape: n_cases * n_out
        '''
        # Append 1 to each list of input vals, for the constant term:
        x_plus = torch.cat([x, torch.ones(x.shape[0], 1)], dim=1)
        # Shape of weighted_x is n_cases * n_rules * (n_in+1)
        weighted_x = torch.einsum('bp, bq -> bpq', weights, x_plus)
        #print('w,xp:', weights, x_plus)
        #print('wX:', weighted_x)
        # Can't have value 0 for weights, or LSE won't work:
        weighted_x[weighted_x == 0] = 1e-12
        # Squash x and y down to 2D matrices for gels:
        weighted_x_2d = weighted_x.view(weighted_x.shape[0], -1)
        y_actual_2d = y_actual.view(y_actual.shape[0], -1)
        #print('y_actual_2d, weighted_x_2d', y_actual_2d, weighted_x_2d)
        # Use gels to do LSE, then pick out the solution rows:
        try:
        #    coeff_2d, _ = torch.gels(y_actual_2d, weighted_x_2d)
        #except AttributeError:
            coeff_2d, _ = torch.lstsq(y_actual_2d, weighted_x_2d)
        except RuntimeError as e:
            print('Internal error in gels', e)
            print('Weights are:', weighted_x)
            raise e
        coeff_2d = coeff_2d[0:weighted_x_2d.shape[1]]
        # Reshape to 3D tensor: divide by rules, n_in+1, then swap last 2 dims
        self.coeff = coeff_2d.view(weights.shape[1], x.shape[1] + 1, -1) \
            .transpose(1, 2)
//  [python]
def fit_coeff(self, x, weights, y_actual):
        '''
            Use LSE to solve for coeff: y_actual = coeff * (weighted)x
                  x.shape: n_cases * n_in
            weights.shape: n_cases * n_rules
            [ coeff.shape: n_rules * n_out * (n_in+1) ]
                  y.shape: n_cases * n_out
        '''
        # Append 1 to each list of input vals, for the constant term:
        x_plus = torch.cat([x, torch.ones(x.shape[0], 1)], dim=1)
        # Shape of weighted_x is n_cases * n_rules * (n_in+1)
        weighted_x = torch.einsum('bp, bq -> bpq', weights, x_plus)
        #print('w,xp:', weights, x_plus)
        #print('wX:', weighted_x)
        # Can't have value 0 for weights, or LSE won't work:
        weighted_x[weighted_x == 0] = 1e-12
        # Squash x and y down to 2D matrices for gels:
        weighted_x_2d = weighted_x.view(weighted_x.shape[0], -1)
        y_actual_2d = y_actual.view(y_actual.shape[0], -1)
        #print('y_actual_2d, weighted_x_2d', y_actual_2d, weighted_x_2d)
        # Use gels to do LSE, then pick out the solution rows:
        try:
            coeff_2d, _ = torch.linalg.lstsq(y_actual_2d, weighted_x_2d)
        except RuntimeError as e:
            print('Internal error in gels', e)
            print('Weights are:', weighted_x)
            raise e
        coeff_2d = coeff_2d[0:weighted_x_2d.shape[1]]
        # Reshape to 3D tensor: divide by rules, n_in+1, then swap last 2 dims
        self.coeff = coeff_2d.view(weights.shape[1], x.shape[1] + 1, -1) \
            .transpose(1, 2)
//  [python]
def test_scalar_check(self, device):
        zero_d = torch.randn((), device=device)
        one_d = torch.randn((1,), device=device)

        # _multinomial_alias_setup
        self.assertRaises(RuntimeError, lambda: torch._multinomial_alias_setup(zero_d))

        # remainder
        self.assertEqual((), torch.remainder(zero_d, zero_d).shape)
        self.assertEqual((), torch.remainder(zero_d, 2).shape)
        self.assertEqual((1,), torch.remainder(zero_d, one_d).shape)
        self.assertEqual((1,), torch.remainder(one_d, zero_d).shape)

        # fmod
        self.assertEqual((), torch.fmod(zero_d, zero_d).shape)
        self.assertEqual((), torch.fmod(zero_d, 2).shape)
        self.assertEqual((1,), torch.fmod(zero_d, one_d).shape)
        self.assertEqual((1,), torch.fmod(one_d, zero_d).shape)

        # cumsum / cumprod
        self.assertEqual((), torch.cumsum(zero_d, 0).shape)
        self.assertEqual((), torch.cumprod(zero_d, 0).shape)

        # renorm
        self.assertRaises(RuntimeError, lambda: torch.renorm(zero_d, 0.5, 0, 1.0))

        # sort
        self.assertEqual([(), ()], [x.shape for x in torch.sort(zero_d, 0, False)])
        self.assertEqual([(), ()], [x.shape for x in torch.sort(zero_d, 0, True)])

        # lstsq (gels)
        self.assertRaises(RuntimeError, lambda: torch.lstsq(zero_d, zero_d))

        # eig
        self.assertRaises(RuntimeError, lambda: torch.eig(zero_d, False))
        self.assertRaises(RuntimeError, lambda: torch.eig(zero_d, True))

        # max, min
        self.assertEqual((), torch.max(zero_d, zero_d).shape)
        self.assertEqual((1,), torch.max(one_d, zero_d).shape)
        self.assertEqual((1,), torch.max(zero_d, one_d).shape)
        self.assertEqual((), torch.min(zero_d, zero_d).shape)
        self.assertEqual((1,), torch.min(one_d, zero_d).shape)
        self.assertEqual((1,), torch.min(zero_d, one_d).shape)

        # diag
        self.assertRaises(RuntimeError, lambda: torch.diag(zero_d))
//  [python]
def test_scalar_check(self, device):
        zero_d = torch.randn((), device=device)
        one_d = torch.randn((1,), device=device)

        # _multinomial_alias_setup
        self.assertRaises(RuntimeError, lambda: torch._multinomial_alias_setup(zero_d))

        # remainder
        self.assertEqual((), torch.remainder(zero_d, zero_d).shape)
        self.assertEqual((), torch.remainder(zero_d, 2).shape)
        self.assertEqual((1,), torch.remainder(zero_d, one_d).shape)
        self.assertEqual((1,), torch.remainder(one_d, zero_d).shape)

        # fmod
        self.assertEqual((), torch.fmod(zero_d, zero_d).shape)
        self.assertEqual((), torch.fmod(zero_d, 2).shape)
        self.assertEqual((1,), torch.fmod(zero_d, one_d).shape)
        self.assertEqual((1,), torch.fmod(one_d, zero_d).shape)

        # cumsum / cumprod
        self.assertEqual((), torch.cumsum(zero_d, 0).shape)
        self.assertEqual((), torch.cumprod(zero_d, 0).shape)

        # renorm
        self.assertRaises(RuntimeError, lambda: torch.renorm(zero_d, 0.5, 0, 1.0))

        # sort
        self.assertEqual([(), ()], [x.shape for x in torch.sort(zero_d, 0, False)])
        self.assertEqual([(), ()], [x.shape for x in torch.sort(zero_d, 0, True)])

        # lstsq (gels)
        self.assertRaises(RuntimeError, lambda: torch.linalg.lstsq(zero_d, zero_d))

        # eig
        self.assertRaises(RuntimeError, lambda: torch.eig(zero_d, False))
        self.assertRaises(RuntimeError, lambda: torch.eig(zero_d, True))

        # max, min
        self.assertEqual((), torch.max(zero_d, zero_d).shape)
        self.assertEqual((1,), torch.max(one_d, zero_d).shape)
        self.assertEqual((1,), torch.max(zero_d, one_d).shape)
        self.assertEqual((), torch.min(zero_d, zero_d).shape)
        self.assertEqual((1,), torch.min(one_d, zero_d).shape)
        self.assertEqual((1,), torch.min(zero_d, one_d).shape)

        # diag
        self.assertRaises(RuntimeError, lambda: torch.diag(zero_d))
//  [python]
def _fit_2d_polynomial_surface(
    coords: Tensor, values: Tensor, order: int, surface_shape: Size
) -> Tensor:
    """Fit a 2D polynomial to a set of coordinates and their values,
    and return the surface evaluated at every point."""
    n_coeffs = int((order + 1) * (order + 2) / 2)
    if n_coeffs >= len(values):
        raise ValueError(
            f"Cannot fit a {order} degree 2D polynomial "
            f"with {len(values)} sampled values"
        )
    orders = torch.arange(order + 1, device=coords.device)
    order_pairs = torch.stack(torch.meshgrid(orders, orders), -1)
    order_pairs = order_pairs[order_pairs.sum(-1) <= order].reshape(-1, 2)
    terms = torch.stack(
        [coords[:, 0] ** i * coords[:, 1] ** j for i, j in order_pairs], -1
    )
    # use "gels" driver for precision and GPU consistency
    coeffs = torch.lstsq(terms, values)[0]
    dense_coords = torch.meshgrid(
        [
            torch.arange(s, dtype=values.dtype, device=values.device)
            for s in surface_shape
        ]
    )
    dense_terms = torch.stack(
        [dense_coords[0] ** i * dense_coords[1] ** j for i, j in order_pairs],
        -1,
    )
    return torch.matmul(dense_terms, coeffs)
//  [python]
def _fit_2d_polynomial_surface(
    coords: Tensor, values: Tensor, order: int, surface_shape: Size
) -> Tensor:
    """Fit a 2D polynomial to a set of coordinates and their values,
    and return the surface evaluated at every point."""
    n_coeffs = int((order + 1) * (order + 2) / 2)
    if n_coeffs >= len(values):
        raise ValueError(
            f"Cannot fit a {order} degree 2D polynomial "
            f"with {len(values)} sampled values"
        )
    orders = torch.arange(order + 1, device=coords.device)
    order_pairs = torch.stack(torch.meshgrid(orders, orders), -1)
    order_pairs = order_pairs[order_pairs.sum(-1) <= order].reshape(-1, 2)
    terms = torch.stack(
        [coords[:, 0] ** i * coords[:, 1] ** j for i, j in order_pairs], -1
    )
    # use "gels" driver for precision and GPU consistency
    coeffs = torch.linalg.lstsq(terms, values, driver="gels").solution
    dense_coords = torch.meshgrid(
        [
            torch.arange(s, dtype=values.dtype, device=values.device)
            for s in surface_shape
        ]
    )
    dense_terms = torch.stack(
        [dense_coords[0] ** i * dense_coords[1] ** j for i, j in order_pairs],
        -1,
    )
    return torch.matmul(dense_terms, coeffs)
//  [python]
def __call__(self, xs, ys, inds=None):
        xs, ys = xs.cpu(), ys.cpu()
        if inds is None:
            inds = range(ys.shape[1])
        else:
            if max(inds) >= ys.shape[1] or min(inds) < 0:
                raise ValueError("inds contain indices where xs and ys are not defined")

        preds = []

        for i in inds:
            if i == 0:
                preds.append(torch.zeros_like(ys[:, 0]))  # predict zero for first point
                continue
            train_xs, train_ys = xs[:, :i], ys[:, :i]
            test_x = xs[:, i : i + 1]

            ### weighting decaying covariance
            # Sigma = torch.tensor([i**(-2) for i in range(1, 20+1)])
            # # Sigma = torch.cat([torch.ones(3),0.05*torch.ones(20-3)])
            # Sigma = torch.sqrt(torch.diag(Sigma))
            # train_xs = train_xs @ Sigma
            ws, _, _, _ = torch.lstsq(
                train_xs, train_ys.unsqueeze(2)
            )
            # ws = Sigma@ws

            pred = test_x @ ws
            preds.append(pred[:, 0, 0])

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

        preds = []

        for i in inds:
            if i == 0:
                preds.append(torch.zeros_like(ys[:, 0]))  # predict zero for first point
                continue
            train_xs, train_ys = xs[:, :i], ys[:, :i]
            test_x = xs[:, i : i + 1]

            ### weighting decaying covariance
            # Sigma = torch.tensor([i**(-2) for i in range(1, 20+1)])
            # # Sigma = torch.cat([torch.ones(3),0.05*torch.ones(20-3)])
            # Sigma = torch.sqrt(torch.diag(Sigma))
            # train_xs = train_xs @ Sigma
            ws, _, _, _ = torch.linalg.lstsq(
                train_xs, train_ys.unsqueeze(2), driver=self.driver
            )
            # ws = Sigma@ws

            pred = test_x @ ws
            preds.append(pred[:, 0, 0])

        return torch.stack(preds, dim=1)
//  [python]
def fit_coeff(self, x, weights, y_actual):
        '''
            Use LSE to solve for coeff: y_actual = coeff * (weighted)x
                  x.shape: n_cases * n_in
            weights.shape: n_cases * n_rules
            [ coeff.shape: n_rules * n_out * (n_in+1) ]
                  y.shape: n_cases * n_out
        '''
        # Append 1 to each list of input vals, for the constant term:
        #x_plus = torch.cat([x, torch.ones(x.shape[0], 1)], dim=1)

        x_plus = torch.ones(len(x), 1)
        #x_plus = torch.cat([x, torch.ones(x.shape[0], 1)], dim=1)

        # Shape of weighted_x is n_cases * n_rules * (n_in+1)
        weighted_x = torch.einsum('bp, bq -> bpq', weights, x_plus)

        # Can't have value 0 for weights, or LSE won't work:
        weighted_x[weighted_x == 0] = 1e-12
        # Squash x and y down to 2D matrices for gels:
        weighted_x_2d = weighted_x.view(weighted_x.shape[0], -1)
        y_actual_2d = y_actual.view(y_actual.shape[0], -1)

        # Use gels to do LSE, then pick out the solution rows:

        try:
            #lstqs usato nella nuova versione di Pytorch
            coeff_2d, _ = torch.lstsq(y_actual_2d, weighted_x_2d) #torch.gels(y_actual_2d, weighted_x_2d)

        except RuntimeError as e:
            print('Internal error in gels', e)
            print('Weights are:', weighted_x)
            raise e

        coeff_2d = coeff_2d[0:weighted_x_2d.shape[1]]


        # Reshape to 3D tensor: divide by rules, n_in+1, then swap last 2 dims

        #self.coeff = coeff_2d.view(weights.shape[1], x.shape[1]+1, -1).transpose(1, 2)

        self.coeff = coeff_2d.view(weights.shape[1], self.d_out, -1).transpose(1, 2)
//  [python]
def fit_coeff(self, x, weights, y_actual):
        '''
            Use LSE to solve for coeff: y_actual = coeff * (weighted)x
                  x.shape: n_cases * n_in
            weights.shape: n_cases * n_rules
            [ coeff.shape: n_rules * n_out * (n_in+1) ]
                  y.shape: n_cases * n_out
        '''
        # Append 1 to each list of input vals, for the constant term:
        #x_plus = torch.cat([x, torch.ones(x.shape[0], 1)], dim=1)

        x_plus = torch.ones(len(x), 1)
        #x_plus = torch.cat([x, torch.ones(x.shape[0], 1)], dim=1)

        # Shape of weighted_x is n_cases * n_rules * (n_in+1)
        weighted_x = torch.einsum('bp, bq -> bpq', weights, x_plus)

        # Can't have value 0 for weights, or LSE won't work:
        weighted_x[weighted_x == 0] = 1e-12
        # Squash x and y down to 2D matrices for gels:
        weighted_x_2d = weighted_x.view(weighted_x.shape[0], -1)
        y_actual_2d = y_actual.view(y_actual.shape[0], -1)

        # Use gels to do LSE, then pick out the solution rows:

        try:
            #lstqs usato nella nuova versione di Pytorch
            coeff_2d, _ = torch.linalg.lstsq(y_actual_2d, weighted_x_2d) #torch.gels(y_actual_2d, weighted_x_2d)

        except RuntimeError as e:
            print('Internal error in gels', e)
            print('Weights are:', weighted_x)
            raise e

        coeff_2d = coeff_2d[0:weighted_x_2d.shape[1]]


        # Reshape to 3D tensor: divide by rules, n_in+1, then swap last 2 dims

        #self.coeff = coeff_2d.view(weights.shape[1], x.shape[1]+1, -1).transpose(1, 2)

        self.coeff = coeff_2d.view(weights.shape[1], self.d_out, -1).transpose(1, 2)
//  [python]
def fit_coeff(self, x, weights, y_actual):
        '''
            Use LSE to solve for coeff: y_actual = coeff * (weighted)x
                  x.shape: n_cases * n_in
            weights.shape: n_cases * n_rules
            [ coeff.shape: n_rules * n_out * (n_in+1) ]
                  y.shape: n_cases * n_out
        '''
        # Append 1 to each list of input vals, for the constant term:
        #x_plus = torch.cat([x, torch.ones(x.shape[0], 1)], dim=1)

        x_plus = torch.ones(len(x), 1)

        # Shape of weighted_x is n_cases * n_rules * (n_in+1)
        weighted_x = torch.einsum('bp, bq -> bpq', weights, x_plus)

        # Can't have value 0 for weights, or LSE won't work:
        weighted_x[weighted_x == 0] = 1e-12
        # Squash x and y down to 2D matrices for gels:
        weighted_x_2d = weighted_x.view(weighted_x.shape[0], -1)
        y_actual_2d = y_actual.view(y_actual.shape[0], -1)

        #print('weighted_x_2d',weighted_x_2d) #PESI
        #print('y_actual_2d',y_actual_2d) #RAPPRESENTAZIONE ONE OUT DEL TARGET

        # Use gels to do LSE, then pick out the solution rows:
        try:
            #coeff_2d, _ = torch.gels(y_actual_2d, weighted_x_2d)
            coeff_2d, _ = torch.lstsq(y_actual_2d, weighted_x_2d)
        except RuntimeError as e:
            print('Internal error in gels', e)
            print('Weights are:', weighted_x)
            raise e
        coeff_2d = coeff_2d[0:weighted_x_2d.shape[1]]
        # Reshape to 3D tensor: divide by rules, n_in+1, then swap last 2 dims

        #self.coeff = coeff_2d.view(weights.shape[1], x.shape[1]+1, -1).transpose(1, 2)

        self.coeff = coeff_2d.view(weights.shape[1], self.d_out, -1).transpose(1, 2)