jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def infer_meta_from_obs_action_and_rewards(self, obs: torch.Tensor, action: 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(torch.cat([obs, action], dim=1))
        z = torch.solve(torch.mv(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_action_and_rewards(self, obs: torch.Tensor, action: 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(torch.cat([obs, action], dim=1))
        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 derive_matrix_from_vertexes(init_vertex, final_vertex):
    '''
    init_vertex, final_vertex: tensor of shape (B,2,H,W) or (B,N,2)
    '''
    init_vertex  = init_vertex.flatten(2)
    final_vertex = final_vertex.flatten(2)
    if init_vertex.shape[1] != 2:
        init_vertex = init_vertex.permute(0,2,1)
    if final_vertex.shape[1] != 2:
        final_vertex = final_vertex.permute(0,2,1)

    X, Y = init_vertex[:,0], init_vertex[:,1]
    U, V = final_vertex[:,0], final_vertex[:,1]
    O, I = torch.zeros_like(X), torch.ones_like(X)
    A = torch.cat([torch.stack([X, Y, I, O, O, O, -U*X, -U*Y, O], dim=-1),
                   torch.stack([O, O, O, X, Y, I, -V*X, -V*Y, O], dim=-1)], dim=1)
    B = torch.cat([U,V],dim=1)
    X = torch.solve(B,A).solution
    # X = torch.linalg.lstsq(A, B).solution
    I = torch.ones(size=(X.shape[0],1), dtype=X.dtype, device=X.device)
    X = torch.cat([X[:,:8], I], dim=1).reshape(X.shape[0], 3, -1)
    return X
//  [python]
def derive_matrix_from_vertexes(init_vertex, final_vertex):
    '''
    init_vertex, final_vertex: tensor of shape (B,2,H,W) or (B,N,2)
    '''
    init_vertex  = init_vertex.flatten(2)
    final_vertex = final_vertex.flatten(2)
    if init_vertex.shape[1] != 2:
        init_vertex = init_vertex.permute(0,2,1)
    if final_vertex.shape[1] != 2:
        final_vertex = final_vertex.permute(0,2,1)

    X, Y = init_vertex[:,0], init_vertex[:,1]
    U, V = final_vertex[:,0], final_vertex[:,1]
    O, I = torch.zeros_like(X), torch.ones_like(X)
    A = torch.cat([torch.stack([X, Y, I, O, O, O, -U*X, -U*Y, O], dim=-1),
                   torch.stack([O, O, O, X, Y, I, -V*X, -V*Y, O], dim=-1)], dim=1)
    B = torch.cat([U,V],dim=1)
    X = torch.linalg.lstsq(A,B,driver='gels').solution
    # X = torch.linalg.lstsq(A, B).solution
    I = torch.ones(size=(X.shape[0],1), dtype=X.dtype, device=X.device)
    X = torch.cat([X[:,:8], I], dim=1).reshape(X.shape[0], 3, -1)
    return X
//  [python]
def project_function_sparse(
    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[None]
    )  # 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.solve(basis_vals.t() @ func_value, basis_vals.t() @ basis_vals).solution
    soln = soln.T.reshape(batch_size, -1)
    others = others[:, :1, :]
    others = others.reshape(batch_size, -1)
    return soln, others
//  [python]
def project_function_sparse(
    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[None]
    )  # 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(basis_vals, func_value).solution
    soln = soln.T.reshape(batch_size, -1)
    others = others[:, :1, :]
    others = others.reshape(batch_size, -1)
    return soln, others
//  [python]
def find_best_perm_with_least_square(est_targets, targets):
        r"""Find the best permutation via the least square.

        Args:
            est_targets: torch.Tensor, shape (..., n_sources, n_samples).
                Estimated separation outputs in time-domain.
            targets: torch.Tensor, shape (..., n_mixtures, n_samples).
                Mixtures (used as ground truths) in time-domain.

        Returns:
            perm_mat: torch.Tensor, shape (..., n_sources, n_mixtures).
                Optimal permutation matrix.
        """
        # (..., n_samples, n_sources)
        est_targets = est_targets.transpose(-1, -2)
        # (..., n_samples, n_mixtures)
        targets = targets.transpose(-1, -2)

        # least-square
        A = torch.solve(targets, est_targets)[0]
        # get maximum number to obtain permutation matrix
        Amax, _ = torch.max(A, dim=-1, keepdim=True)
        # matrix whose component is 1 if the element is maximum number else 0
        perm_mat = (A == Amax).to(est_targets.dtype)

        return perm_mat.transpose(-1, -2)
//  [python]
def find_best_perm_with_least_square(est_targets, targets):
        r"""Find the best permutation via the least square.

        Args:
            est_targets: torch.Tensor, shape (..., n_sources, n_samples).
                Estimated separation outputs in time-domain.
            targets: torch.Tensor, shape (..., n_mixtures, n_samples).
                Mixtures (used as ground truths) in time-domain.

        Returns:
            perm_mat: torch.Tensor, shape (..., n_sources, n_mixtures).
                Optimal permutation matrix.
        """
        # (..., n_samples, n_sources)
        est_targets = est_targets.transpose(-1, -2)
        # (..., n_samples, n_mixtures)
        targets = targets.transpose(-1, -2)

        # least-square
        A = torch.linalg.lstsq(est_targets, targets)[0]
        # get maximum number to obtain permutation matrix
        Amax, _ = torch.max(A, dim=-1, keepdim=True)
        # matrix whose component is 1 if the element is maximum number else 0
        perm_mat = (A == Amax).to(est_targets.dtype)

        return perm_mat.transpose(-1, -2)
//  [python]
def _get_average_E0(self) -> Optional[Dict[int, float]]:
        """
        Function to compute the average interaction energy of each chemical element
        returns dictionary of E0s
        """
        if self.atomic_energies == "auto":
            numbers = [atomic_numbers[s] for s in self._get_species()]
            num_elements = len(numbers)
            len_train = len(self.train_dataset)
            A = torch.zeros((len_train, num_elements))
            B = torch.zeros((len_train,))
            
            for i, sample in enumerate(self.train_dataset):
                B[i] = sample[properties.energy]
                for j, z in enumerate(numbers):
                    A[i, j] = torch.count_nonzero(sample[properties.Z] == z)
            atomic_energies_dict = {z: 0.0 for z in numbers}
            
            try:
                E0s = torch.solve(torch.matmul(A.t(),B),torch.matmul(A.t(),A))[0]
                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"
                )
            self.atomic_energies = atomic_energies_dict
            
        elif isinstance(self.atomic_energies, Dict):
            _atomic_energies = {}
            for k, v in self.atomic_energies.items():
                if isinstance(k, int):
                    _atomic_energies[chemical_symbols[k]] = v
            self.atomic_energies = _atomic_energies

        logger.info(f"Using reference energies for elements: {self.atomic_energies}.")
        return self.atomic_energies
//  [python]
def _get_average_E0(self) -> Optional[Dict[int, float]]:
        """
        Function to compute the average interaction energy of each chemical element
        returns dictionary of E0s
        """
        if self.atomic_energies == "auto":
            numbers = [atomic_numbers[s] for s in self._get_species()]
            num_elements = len(numbers)
            len_train = len(self.train_dataset)
            A = torch.zeros((len_train, num_elements))
            B = torch.zeros((len_train,))
            
            for i, sample in enumerate(self.train_dataset):
                B[i] = sample[properties.energy]
                for j, z in enumerate(numbers):
                    A[i, j] = torch.count_nonzero(sample[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"
                )
            self.atomic_energies = atomic_energies_dict
            
        elif isinstance(self.atomic_energies, Dict):
            _atomic_energies = {}
            for k, v in self.atomic_energies.items():
                if isinstance(k, int):
                    _atomic_energies[chemical_symbols[k]] = v
            self.atomic_energies = _atomic_energies

        logger.info(f"Using reference energies for elements: {self.atomic_energies}.")
        return self.atomic_energies
//  [python]
def calculate_smoothed_slope(
    losses: list[float],
    smooth_window_size: int = 3,
    smoothing_constant: float = 0.93,
    regression_win: int = 20,
) -> float:
    y = torch.tensor(losses[-smooth_window_size - 10 :]).view(-1, 1)
    x = torch.arange(y.shape[0]).view(-1, 1)

    # Smooth the data using a simple moving average with a smoothing constant
    smooth_window_size = 3
    conv_weights = (smoothing_constant / smooth_window_size) * torch.ones(
        smooth_window_size
    )
    smoothed_y = torch.nn.functional.conv1d(
        y.view(1, 1, -1),
        conv_weights.view(1, 1, -1),
        padding=(smooth_window_size - 1) // 2,
    ).view(-1)

    regression_win = min(regression_win, smoothed_y.shape[-1])
    # Perform linear regression on the smoothed_y
    x_last_n = x[-regression_win:]
    y_last_n = smoothed_y[-regression_win:]
    X = torch.cat([torch.ones(regression_win, 1), x_last_n], dim=1)
    coefficients = torch.solve(X.t() @ y_last_n.view(-1, 1), X.t() @ X).solution.flatten()
    return coefficients[1]
//  [python]
def calculate_smoothed_slope(
    losses: list[float],
    smooth_window_size: int = 3,
    smoothing_constant: float = 0.93,
    regression_win: int = 20,
) -> float:
    y = torch.tensor(losses[-smooth_window_size - 10 :]).view(-1, 1)
    x = torch.arange(y.shape[0]).view(-1, 1)

    # Smooth the data using a simple moving average with a smoothing constant
    smooth_window_size = 3
    conv_weights = (smoothing_constant / smooth_window_size) * torch.ones(
        smooth_window_size
    )
    smoothed_y = torch.nn.functional.conv1d(
        y.view(1, 1, -1),
        conv_weights.view(1, 1, -1),
        padding=(smooth_window_size - 1) // 2,
    ).view(-1)

    regression_win = min(regression_win, smoothed_y.shape[-1])
    # Perform linear regression on the smoothed_y
    x_last_n = x[-regression_win:]
    y_last_n = smoothed_y[-regression_win:]
    X = torch.cat([torch.ones(regression_win, 1), x_last_n], dim=1)
    coefficients = torch.linalg.lstsq(X, y_last_n.view(-1, 1)).solution.flatten()
    return coefficients[1]
//  [python]
def _dls_torch(X: T, ok_mask: T, C: T, Y: T, n: int, cuda_support: bool) -> None:
    for i in range(n):
        mask_i = ok_mask[:, i]
        if cuda_support:
            try:
                Y[:, i] = torch.solve(C[mask_i].t() @ C[mask_i], C[mask_i].t() @ X[mask_i, i])[0]
            except RuntimeError:
                print("Falling back")
                Y[:, i] = torch.pinverse(C[mask_i]) @ X[mask_i, i]
        else:
            try:
                Y[:, i] = torch.solve(C[mask_i].t() @ C[mask_i], C[mask_i].t() @ X[mask_i, i])[0]
            except RuntimeError:
                print("Falling back")
                Y[:, i] = torch.solve(C[mask_i].t() @ C[mask_i], C[mask_i].t() @ X[mask_i, i])[0]
//  [python]
def _dls_torch(X: T, ok_mask: T, C: T, Y: T, n: int, cuda_support: bool) -> None:
    for i in range(n):
        mask_i = ok_mask[:, i]
        if cuda_support:
            try:
                Y[:, i] = torch.linalg.lstsq(C[mask_i], X[mask_i, i], driver="gels")[0]
            except RuntimeError:
                print("Falling back")
                Y[:, i] = torch.linalg.pinv(C[mask_i]) @ X[mask_i, i]
        else:
            try:
                Y[:, i] = torch.linalg.lstsq(C[mask_i], X[mask_i, i], driver="gelsy")[0]
            except RuntimeError:
                print("Falling back")
                Y[:, i] = torch.linalg.lstsq(C[mask_i], X[mask_i, i], driver="gelsd")[0]
//  [python]
def _solve_dLdZ(self, dLdW):
        """Generates noisy dLdP using de-aggregation trick"""
        A = self.last_input
        B = dLdW

        # Apply pseudoinverse
        dLdZ = B.matmul(A.pinverse()).t()
        return dLdZ
//  [python]
def _solve_dLdZ(self, dLdW):
        """Generates noisy dLdP using de-aggregation trick"""
        A = self.last_input
        B = dLdW

        # Apply pseudoinverse
        dLdZ = torch.linalg.lstsq(A.t(), B.t()).solution
        # dLdZ = B.matmul(A.pinverse()).t()
        return dLdZ
//  [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 @ A) @ A.T # AX=B
        return X.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 solve_output(self, feature, target):

        ftf = feature.T @ feature
        fty = feature.T @ target

        try:
            output_w = torch.solve(fty, ftf)[0].to(self.opts.device)
        except:
            output_w = torch.mm(torch.pinverse(ftf), fty).to(self.opts.device)
            
        return output_w
//  [python]
def solve_output(self, feature, target):

        ftf = feature.T @ feature
        fty = feature.T @ target

        try:
            output_w = torch.linalg.lstsq(ftf, fty, driver = 'gelsd').solution.to(self.opts.device)
        except:
            output_w = torch.mm(torch.linalg.pinv(ftf), fty).to(self.opts.device)
            
        return output_w
//  [python]
def get_depth_anchors(traj_gt, num_obs, num_full, offset_anchor='Zc', time_orders=None):
    if offset_anchor == 'Zc':
        anchors = traj_gt[num_obs-1: num_obs, 2:3]
    else:
        zs_obs = traj_gt[:num_obs, 2:3]
        if isinstance(zs_obs, np.ndarray):
            zs_obs = torch.from_numpy(zs_obs).to(time_orders.device)
        solution, _ = torch.solve(zs_obs, time_orders[:num_obs, :])
        anchors = torch.matmul(time_orders[num_obs:num_full, :], solution)
        if isinstance(traj_gt, np.ndarray):
            anchors = anchors.cpu().numpy()
    return anchors
//  [python]
def get_depth_anchors(traj_gt, num_obs, num_full, offset_anchor='Zc', time_orders=None):
    if offset_anchor == 'Zc':
        # use the last observed depth as the offset anchor
        anchors = traj_gt[num_obs-1: num_obs, 2:3]
    else:
        zs_obs = traj_gt[:num_obs, 2:3]
        if isinstance(zs_obs, np.ndarray):
            zs_obs = torch.from_numpy(zs_obs).to(time_orders.device)
        # use historical (z, t) values to fit a four-order multinomial model
        solution = torch.linalg.lstsq(time_orders[:num_obs, :], zs_obs).solution  # (5,1)
        anchors = torch.matmul(time_orders[num_obs:num_full, :], solution)
        if isinstance(traj_gt, np.ndarray):
            anchors = anchors.cpu().numpy()
    return anchors
//  [python]
def polynomial_fit(self, R, P, V):
        A = self.evaluate(R, P, require_matrix=True)
        B = V.reshape(-1)
        X, residuals, rank, singular_values = torch.lstsq(A, B)
        print("fitted coefficients:{}".format(X))
        print("lstsq residuals:{}".format(residuals))
        self.c = X
//  [python]
def polynomial_fit(self, R, P, V):
        A = self.evaluate(R, P, require_matrix=True)
        B = V.reshape(-1)
        X, residuals, rank, singular_values = torch.linalg.lstsq(A, B, driver='gels')
        print("fitted coefficients:{}".format(X))
        print("lstsq residuals:{}".format(residuals))
        self.c = X
//  [python]
def getTrueVariationalMean(t, latent_mean, inducing_points_locs, kernel):
    Ktz = kernel.buildKernelMatrix(X1=t, X2=inducing_points_locs)
    Kzz = kernel.buildKernelMatrix(X1=inducing_points_locs, X2=inducing_points_locs)
    res = torch.lstsq(latent_mean, Ktz)
    vMean = torch.matmul(Kzz, res.solution)
    return vMean
//  [python]
def getTrueVariationalMean(t, latent_mean, inducing_points_locs, kernel):
    Ktz = kernel.buildKernelMatrix(X1=t, X2=inducing_points_locs)
    Kzz = kernel.buildKernelMatrix(X1=inducing_points_locs, X2=inducing_points_locs)
    res = torch.linalg.lstsq(Ktz, latent_mean)
    vMean = torch.matmul(Kzz, res.solution)
    return vMean
//  [python]
def affine_solver(
    source_pts: torch.tensor, target_pts: torch.tensor
) -> Tuple[torch.tensor, torch.tensor]:
    """
    Compute the affine transform and inverse affine transform with given source and target points.
    Currently, this function only supports points in 3D. Noted, this function does not support batch computation.
    Args:
        source_pts (torch.tensor): Nx3
        target_pts (torch.tensor): Nx3

    Returns:
        torch.tensor: The 3x4 affine matrix.
        torch.tensor: The 3x4 inverse affine matrix.

    """

    affine_matrix = torch.solve(target_pts, source_pts).solution
    affine_matrix_inv = torch.inverse(affine_matrix.transpose(0, 1))
    return affine_matrix.transpose(0, 1)[:3], affine_matrix_inv[:3]
//  [python]
def affine_solver(
    source_pts: torch.tensor, target_pts: torch.tensor
) -> Tuple[torch.tensor, torch.tensor]:
    """
    Compute the affine transform and inverse affine transform with given source and target points.
    Currently, this function only supports points in 3D. Noted, this function does not support batch computation.
    Args:
        source_pts (torch.tensor): Nx3
        target_pts (torch.tensor): Nx3

    Returns:
        torch.tensor: The 3x4 affine matrix.
        torch.tensor: The 3x4 inverse affine matrix.

    """

    affine_matrix = torch.linalg.lstsq(source_pts, target_pts).solution
    affine_matrix_inv = torch.linalg.inv(affine_matrix.transpose(0, 1))
    return affine_matrix.transpose(0, 1)[:3], affine_matrix_inv[:3]
//  [python]
def _loglik(Sigma: torch.Tensor, Sigma_hat: torch.Tensor, n: int) -> float:
    p = Sigma.shape[0]
    l = (n/2)*(p*np.log(2*np.pi) +
               torch.logdet(Sigma) +
               torch.trace(torch.solve(Sigma_hat, Sigma)[1]))
    return l
//  [python]
def _loglik(Sigma: torch.Tensor, Sigma_hat: torch.Tensor, n: int) -> float:
    p = Sigma.shape[0]
    l = (n/2)*(p*np.log(2*np.pi) +
               torch.logdet(Sigma) +
               torch.trace(torch.linalg.lstsq(Sigma, Sigma_hat).solution))
    return l
//  [python]
def calibrate_probs(probs, batch, model):
    Diags = get_calibrate_matrices(batch,model)
    for i in range(probs.shape[0]):
        probs[i] = torch.lstsq(probs[i], Diags[i])[0]
        # probs[i] += torch.lstsq(probs[i], Diags[i])[0]
    return probs
//  [python]
def calibrate_probs(probs, batch, model):
    Diags = get_calibrate_matrices(batch,model)
    for i in range(probs.shape[0]):
        probs[i] = torch.linalg.lstsq(Diags[i],probs[i]).solution
        # probs[i] += torch.linalg.lstsq(Diags[i],probs[i]).solution
    return probs
//  [python]
def get_atomref(dataset, prop_name, data_len = None, atomic_number_max = 60):
    '''
    prop_name: "energy"
    '''
    data_len = len(dataset) if data_len is None else data_len
    r = torch.zeros(data_len, atomic_number_max+1)
    energy = torch.zeros(data_len)
    for i in range(data_len):
        data = dataset[i]
        counter = scatter(torch.ones(data.atomic_numbers.shape[0]),data.atomic_numbers.reshape(-1),dim = 0)
        energy[i] = data[prop_name].item()
        r[i,:len(counter)] = counter

    mask = torch.sum(r, dim = 0) > 0
    r = r[:, mask]
    legal_atomref = torch.lstsq(energy, r).solution
    print('in this dataset, available legal atomic numer is {}. its atomref is {}'.format(torch.where(mask)[0],legal_atomref))
    atomref = torch.zeros(atomic_number_max+1)
    atomref[mask] = legal_atomref
    return atomref
//  [python]
def get_atomref(dataset, prop_name, data_len = None, atomic_number_max = 60):
    '''
    prop_name: "energy"
    '''
    data_len = len(dataset) if data_len is None else data_len
    r = torch.zeros(data_len, atomic_number_max+1)
    energy = torch.zeros(data_len)
    for i in range(data_len):
        data = dataset[i]
        counter = scatter(torch.ones(data.atomic_numbers.shape[0]),data.atomic_numbers.reshape(-1),dim = 0)
        energy[i] = data[prop_name].item()
        r[i,:len(counter)] = counter

    mask = torch.sum(r, dim = 0) > 0
    r = r[:, mask]
    legal_atomref = torch.linalg.lstsq(r, energy).solution
    print('in this dataset, available legal atomic numer is {}. its atomref is {}'.format(torch.where(mask)[0],legal_atomref))
    atomref = torch.zeros(atomic_number_max+1)
    atomref[mask] = legal_atomref
    return atomref
//  [python]
def _get_latent_worker(W: torch.tensor, L: torch.tensor, Phi: torch.tensor,
                       Y: torch.tensor, device='cpu', rcond: float = 1e-08):
    if L is not None:
        S21 = torch.cat([W, L], axis=1).to(device)
        S22 = W @ W.T + L @ L.T + Phi
    else:
        S21 = W
        S22 = W @ W.T + Phi
    S22inv_S21 = torch.solve(S21, S22).solution
    E_z_x = Y @ S22inv_S21
    return E_z_x
//  [python]
def _get_latent_worker(W: torch.tensor, L: torch.tensor, Phi: torch.tensor,
                       Y: torch.tensor, device='cpu', rcond: float = 1e-08):
    if L is not None:
        S21 = torch.cat([W, L], axis=1).to(device)
        S22 = W @ W.T + L @ L.T + Phi
    else:
        S21 = W
        S22 = W @ W.T + Phi
    S22inv_S21 = torch.linalg.lstsq(S22, S21, rcond=rcond).solution
    E_z_x = Y @ S22inv_S21
    return E_z_x
//  [python]
def _EM_step_no_private_stable(
        W: torch.Tensor, Phi: torch.Tensor, Y: torch.Tensor,
        Sigma: torch.Tensor, p: List[int], device = 'cpu',
        rcond: float = 1e-08):
    """One EM step when there is no private structure.

    Args:
      W: p_all=sum(p) by d tensor. Current estimate of W.
      Phi: p_all by p_all tensor. Current estimate of Phi.
      Y: N by p_all torch tensor, the data.
      Sigma: p_all by p_all torch tensor, Y.T @ Y / N.
      p: List of integers with the dimensionality of each dataset.
      rcond: Condition number for least squares.
    Returns:
      Tuple W_next, Phi_next with updated parameters.
    """
    # TODO(brielin): use device
    if device == 'gpu': raise NotImplementedError()
    d = W.shape[1]
    N = Y.shape[0]
    Q_inv_W = torch.solve(W @ W.T + Phi, W)[0]
    E_z = Y @ Q_inv_W
    sum_E_zzT = N*torch.eye(d) - N*(W.T @ Q_inv_W) + E_z.T @ E_z
    W_next = torch.solve(sum_E_zzT, (E_z.T @ Y))[0].T
    # Phi_next = Sigma - (Y.T @ E_z @ W_next.T)/N
    Phi_next = Sigma - (W_next @ E_z.T @ Y)/N
    psum = np.concatenate([[0], np.cumsum(p, 0)])
    for i_l, i_r in zip(psum[:-1], psum[1:]):
        Phi_next[i_l:i_r, i_r:] = 0
        Phi_next[i_r:, i_l:i_r] = 0
    return W_next, Phi_next
//  [python]
def _EM_step_no_private_stable(
        W: torch.Tensor, Phi: torch.Tensor, Y: torch.Tensor,
        Sigma: torch.Tensor, p: List[int], device = 'cpu',
        rcond: float = 1e-08):
    """One EM step when there is no private structure.

    Args:
      W: p_all=sum(p) by d tensor. Current estimate of W.
      Phi: p_all by p_all tensor. Current estimate of Phi.
      Y: N by p_all torch tensor, the data.
      Sigma: p_all by p_all torch tensor, Y.T @ Y / N.
      p: List of integers with the dimensionality of each dataset.
      rcond: Condition number for least squares.
    Returns:
      Tuple W_next, Phi_next with updated parameters.
    """
    # TODO(brielin): use device
    if device == 'gpu': raise NotImplementedError()
    d = W.shape[1]
    N = Y.shape[0]
    Q_inv_W = torch.linalg.lstsq(W @ W.T + Phi, W, rcond=rcond).solution
    E_z = Y @ Q_inv_W
    sum_E_zzT = N*torch.eye(d) - N*(W.T @ Q_inv_W) + E_z.T @ E_z
    W_next = torch.linalg.lstsq(sum_E_zzT, (E_z.T @ Y), rcond=rcond).solution.T
    # Phi_next = Sigma - (Y.T @ E_z @ W_next.T)/N
    Phi_next = Sigma - (W_next @ E_z.T @ Y)/N
    psum = np.concatenate([[0], np.cumsum(p, 0)])
    for i_l, i_r in zip(psum[:-1], psum[1:]):
        Phi_next[i_l:i_r, i_r:] = 0
        Phi_next[i_r:, i_l:i_r] = 0
    return W_next, Phi_next
//  [python]
def action_profile_generator_time_ms(args: argparse.Namespace):
    import torch, time
    from . import utils
    from methods.generator import MSGenerator
    
    generator_path: str = args.generator
    tile_count: str = args.tileset
    condition_size: str = args.condition
    batch_size: int = args.batch
    repeat: int = args.repeat
    config = config_tools.get_config_from_namespace(args)

    generator_config = utils.access_yaml(generator_path)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    generator: MSGenerator = generator_config.model_constructor(tile_count, condition_size).to(device)
    
    generation_sizes = config.get("generation_sizes", [])
    
    conditions = torch.rand((batch_size, condition_size), device=device)

    A = torch.ones((len(generation_sizes), 2))
    b = torch.zeros(len(generation_sizes))

    for size_index, size in enumerate(generation_sizes):
        h, w = size
        print(f"{size_index+1}/{len(generation_sizes)}: Working on Size {h}x{w} ...")

        # run once first to ensure initialization is not computed in the time
        generator.generate(conditions, size)

        start = time.time()

        for _ in range(repeat):
            generator.generate(conditions, size)
        
        elapsed = (time.time() - start) / repeat

        A[size_index, 1] = h * w
        b[size_index] = elapsed

        print(f"Generation time / Level = {elapsed} secs")
    
    x, *_ = torch.solve(A.t() @ b, A.t() @ A)

    print(f"Time Estimate = {x[0]} + wh * {x[1]} secs")
    print(f"Correlation Coefficient = { torch.corrcoef(torch.stack([A[:,1], b]))[0, 1] }")
//  [python]
def action_profile_generator_time_ms(args: argparse.Namespace):
    import torch, time
    from . import utils
    from methods.generator import MSGenerator
    
    generator_path: str = args.generator
    tile_count: str = args.tileset
    condition_size: str = args.condition
    batch_size: int = args.batch
    repeat: int = args.repeat
    config = config_tools.get_config_from_namespace(args)

    generator_config = utils.access_yaml(generator_path)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    generator: MSGenerator = generator_config.model_constructor(tile_count, condition_size).to(device)
    
    generation_sizes = config.get("generation_sizes", [])
    
    conditions = torch.rand((batch_size, condition_size), device=device)

    A = torch.ones((len(generation_sizes), 2))
    b = torch.zeros(len(generation_sizes))

    for size_index, size in enumerate(generation_sizes):
        h, w = size
        print(f"{size_index+1}/{len(generation_sizes)}: Working on Size {h}x{w} ...")

        # run once first to ensure initialization is not computed in the time
        generator.generate(conditions, size)

        start = time.time()

        for _ in range(repeat):
            generator.generate(conditions, size)
        
        elapsed = (time.time() - start) / repeat

        A[size_index, 1] = h * w
        b[size_index] = elapsed

        print(f"Generation time / Level = {elapsed} secs")
    
    x, *_ = torch.linalg.lstsq(A, b)

    print(f"Time Estimate = {x[0]} + wh * {x[1]} secs")
    print(f"Correlation Coefficient = { torch.corrcoef(torch.stack([A[:,1], b]))[0, 1] }")
//  [python]
def change_of_basis_proj(x: torch.Tensor, anchors: torch.Tensor) -> torch.Tensor:
    return torch.solve(x.T, anchors.T)[0].T
//  [python]
def change_of_basis_proj(x: torch.Tensor, anchors: torch.Tensor) -> torch.Tensor:
    return torch.linalg.lstsq(anchors.T, x.T)[0].T
//  [python]
def least_important(A, y, rank, size):
    i_glob = torch.zeros(size, dtype=torch.int)
    r_glob = torch.zeros(size, dtype=torch.float)
    i_local = []
    r_local = []
    for i in range(rank, A.size(1), size):
        B = torch.cat([A[:, :i], A[:, i + 1 :]], dim=1)
        mu = torch.lstsq(y, B)[0]
        i_local.append(i)
        r_local.append(cd(B @ mu, y))
    i = torch.argmax(torch.stack(r_local))
    i_glob[rank] = i_local[i]
    r_glob[rank] = r_local[i]
    dist.all_reduce(i_glob)
    dist.all_reduce(r_glob)
    j = torch.argmax(r_glob)
    return i_glob[j], r_glob[j]
//  [python]
def least_important(A, y, rank, size):
    i_glob = torch.zeros(size, dtype=torch.int)
    r_glob = torch.zeros(size, dtype=torch.float)
    i_local = []
    r_local = []
    for i in range(rank, A.size(1), size):
        B = torch.cat([A[:, :i], A[:, i + 1 :]], dim=1)
        mu = torch.linalg.lstsq(B, y).solution
        i_local.append(i)
        r_local.append(cd(B @ mu, y))
    i = torch.argmax(torch.stack(r_local))
    i_glob[rank] = i_local[i]
    r_glob[rank] = r_local[i]
    dist.all_reduce(i_glob)
    dist.all_reduce(r_glob)
    j = torch.argmax(r_glob)
    return i_glob[j], r_glob[j]
//  [python]
def __call__(self, *loss_vals):
        exceeded_iters = False
        if self.max_iters is not None:
            exceeded_iters = self.iters >= self.max_iters

        with torch.no_grad():
            if self.iters % self.frequency == 0 and not exceeded_iters:
                num_losses = self.ratios.shape[-1]
                L = self.W.new(loss_vals[:num_losses])
                if L[L > 0].shape == L.shape:
                    _A = self.A * L * self.W
                    _A = torch.vstack([_A, torch.ones_like(self.W)])

                    # Solve with least squares for weights so each
                    # loss function matches what is given in ratios.
                    X, _, _, _ = torch.lstsq(_A.float(), self.B.float())

                    self.weights = X.tolist()

        self.iters += 1
        return [w * l for w, l in zip(self.weights, loss_vals)]
//  [python]
def __call__(self, *loss_vals):
        exceeded_iters = False
        if self.max_iters is not None:
            exceeded_iters = self.iters >= self.max_iters

        with torch.no_grad():
            if self.iters % self.frequency == 0 and not exceeded_iters:
                num_losses = self.ratios.shape[-1]
                L = self.W.new(loss_vals[:num_losses])
                if L[L > 0].shape == L.shape:
                    _A = self.A * L * self.W
                    _A = torch.vstack([_A, torch.ones_like(self.W)])

                    # Solve with least squares for weights so each
                    # loss function matches what is given in ratios.
                    X = torch.linalg.lstsq(_A.float(), self.B.float(), rcond=None)[0]

                    self.weights = X.tolist()

        self.iters += 1
        return [w * l for w, l in zip(self.weights, loss_vals)]
//  [python]
def scorer_lstsq(oofs, folders, device="cuda:1"):
    data = oofs[folders].values
    target = oofs["target"].values.reshape(-1, 1)

    loo = LeaveOneOut()

    train_X, train_y, valid_X, valid_y = [], [], [], []
    for trn_idx, val_idx in loo.split(data):
        train_X.append(data[trn_idx])
        train_y.append(target[trn_idx])
        valid_X.append(data[val_idx])
        valid_y.append(target[val_idx])

    train_X = torch.tensor(np.stack(train_X), dtype=torch.float32).to(device)
    train_y = torch.tensor(np.stack(train_y), dtype=torch.float32).to(device)
    valid_X = torch.tensor(np.stack(valid_X), dtype=torch.float32).to(device)
    valid_y = torch.tensor(np.stack(valid_y), dtype=torch.float32).to(device)

    W = torch.solve(train_X.T @ train_y, train_X.T @ train_X).solution
    y_pred = valid_X @ W
    mse = torch.nn.functional.mse_loss(y_pred, valid_y)
    return torch.sqrt(mse).cpu().numpy()
//  [python]
def scorer_lstsq(oofs, folders, device="cuda:1"):
    data = oofs[folders].values
    target = oofs["target"].values.reshape(-1, 1)

    loo = LeaveOneOut()

    train_X, train_y, valid_X, valid_y = [], [], [], []
    for trn_idx, val_idx in loo.split(data):
        train_X.append(data[trn_idx])
        train_y.append(target[trn_idx])
        valid_X.append(data[val_idx])
        valid_y.append(target[val_idx])

    train_X = torch.tensor(np.stack(train_X), dtype=torch.float32).to(device)
    train_y = torch.tensor(np.stack(train_y), dtype=torch.float32).to(device)
    valid_X = torch.tensor(np.stack(valid_X), dtype=torch.float32).to(device)
    valid_y = torch.tensor(np.stack(valid_y), dtype=torch.float32).to(device)

    # Same as (np.linalg.pinv(X.T @ X) @ X.T) @ y
    W = torch.linalg.lstsq(train_X, train_y).solution
    y_pred = valid_X @ W
    mse = torch.nn.functional.mse_loss(y_pred, valid_y)
    return torch.sqrt(mse).cpu().numpy()
//  [python]
def _EM_step_full_stable(W: torch.tensor, L: torch.tensor, Phi: torch.tensor,
                         Y: torch.tensor, Sigma: torch.Tensor, p, k,
                         device='cpu', rcond: float = 1e-08):
    """One EM step for the full model.

    Args:
      W: p_all=sum(p) by d tensor. Current estimate of W.
      L: p_all by k_all=sum(k) block diagonal tensor. Current estimate of L.
      Phi: p_all by p_all tensor. Current estimate of Phi.
      Y: N by p_all torch tensor, the data.
      Sigma: p_all by p_all torch tensor, Y.T @ Y / N.
      p:
      k:
      device: 'cpu' or 'gpu'.
      rcond: Condition number for least squares.
    Returns:
      Tuple W_next, L_next, Phi_next
    """
    N, _ = Y.shape
    d = W.shape[1]
    k_all = L.shape[1]
    psum = np.concatenate([[0], np.cumsum(p, 0)])
    ksum = np.concatenate([[0], np.cumsum(k, 0)])

    S21 = torch.cat([W, L], axis=1).to(device)
    S22inv_S21 = torch.solve(S21, W @ W.T + L @ L.T + Phi)[0]
    # TODO(brielin): transpose the next line.
    E_z_x = S22inv_S21.T @ Y.T
    mom_zx_zxT_sum = (N * torch.eye(k_all + d).to(device) - \
                      (N * S22inv_S21.T @ S21).to(device) + \
                      E_z_x @ E_z_x.T).to(device)
    zz_sum = mom_zx_zxT_sum[:d, :d].to(device)
    zx_sum = mom_zx_zxT_sum[:d, d:].to(device)
    xx_sum = mom_zx_zxT_sum[d:, d:].to(device)

    W_next = torch.solve((Y.T @ E_z_x[:d, :].T \
                               - L @ zx_sum.T).T, zz_sum)[0].T
    L_next = torch.solve((Y.T @ E_z_x[d:, :].T \
                               - W @ zx_sum).T, xx_sum)[0].T
    L_next = [L_next[i:j, l:m]
              for i, j, l, m in zip(psum[:-1], psum[1:], ksum[:-1], ksum[1:])]
    L_next = torch.block_diag(*L_next)
    Phi_next = Sigma * N + \
        W @ zz_sum @ W.T + \
        L @ xx_sum @ L.T + \
        2 * L @ zx_sum.T @ W.T - \
        2 * Y.T @ E_z_x[:d, :].T @ W.T - \
        2 * Y.T @ E_z_x[d:, :].T @ L.T
    Phi_next = torch.diag(torch.diagonal(Phi_next / N)).to(device)
    return W_next, L_next, Phi_next
//  [python]
def _EM_step_full_stable(W: torch.tensor, L: torch.tensor, Phi: torch.tensor,
                         Y: torch.tensor, Sigma: torch.Tensor, p, k,
                         device='cpu', rcond: float = 1e-08):
    """One EM step for the full model.

    Args:
      W: p_all=sum(p) by d tensor. Current estimate of W.
      L: p_all by k_all=sum(k) block diagonal tensor. Current estimate of L.
      Phi: p_all by p_all tensor. Current estimate of Phi.
      Y: N by p_all torch tensor, the data.
      Sigma: p_all by p_all torch tensor, Y.T @ Y / N.
      p:
      k:
      device: 'cpu' or 'gpu'.
      rcond: Condition number for least squares.
    Returns:
      Tuple W_next, L_next, Phi_next
    """
    N, _ = Y.shape
    d = W.shape[1]
    k_all = L.shape[1]
    psum = np.concatenate([[0], np.cumsum(p, 0)])
    ksum = np.concatenate([[0], np.cumsum(k, 0)])

    S21 = torch.cat([W, L], axis=1).to(device)
    S22inv_S21 = torch.linalg.lstsq(
        W @ W.T + L @ L.T + Phi, S21, rcond=rcond).solution
    # TODO(brielin): transpose the next line.
    E_z_x = S22inv_S21.T @ Y.T
    mom_zx_zxT_sum = (N * torch.eye(k_all + d).to(device) - \
                      (N * S22inv_S21.T @ S21).to(device) + \
                      E_z_x @ E_z_x.T).to(device)
    zz_sum = mom_zx_zxT_sum[:d, :d].to(device)
    zx_sum = mom_zx_zxT_sum[:d, d:].to(device)
    xx_sum = mom_zx_zxT_sum[d:, d:].to(device)

    W_next = torch.linalg.lstsq(zz_sum, (Y.T @ E_z_x[:d, :].T \
                               - L @ zx_sum.T).T, rcond=rcond).solution.T
    L_next = torch.linalg.lstsq(xx_sum, (Y.T @ E_z_x[d:, :].T \
                               - W @ zx_sum).T, rcond=rcond).solution.T
    L_next = [L_next[i:j, l:m]
              for i, j, l, m in zip(psum[:-1], psum[1:], ksum[:-1], ksum[1:])]
    L_next = torch.block_diag(*L_next)
    Phi_next = Sigma * N + \
        W @ zz_sum @ W.T + \
        L @ xx_sum @ L.T + \
        2 * L @ zx_sum.T @ W.T - \
        2 * Y.T @ E_z_x[:d, :].T @ W.T - \
        2 * Y.T @ E_z_x[d:, :].T @ L.T
    Phi_next = torch.diag(torch.diagonal(Phi_next / N)).to(device)
    return W_next, L_next, Phi_next
//  [python]
def lstsq(input_tensor, A_tensor):
    X = torch.lstsq(input_tensor, A_tensor)
    X[0].requires_grad = True
    X[0].backward(X[0])
    return [X]
//  [python]
def lstsq(input_tensor, A_tensor):
    X = torch.linalg.lstsq(input_tensor, A_tensor)
    X[0].requires_grad = True
    X[0].backward(X[0])
    return [X]
//  [python]
def gradient_descent(
        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]:
    """
    Gradient Descent 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)
    p_list = [p]
    while len(p_list) < max_iter:
        h = -l*torch.solve(torch.matmul(j.T, f).unsqueeze(1), j.T)[0].squeeze()
        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)

        # 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 gradient_descent(
        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]:
    """
    Gradient Descent 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)
    p_list = [p]
    while len(p_list) < max_iter:
        h = -l*torch.linalg.lstsq(j, f, 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)

        # 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 cal_ints_acc(gt_i, pred_i):
    # Red channel:
    gt_i_c = gt_i[:, :1]
    pred_i_c = pred_i[:, :1]
    scale = torch.solve(torch.matmul(pred_i_c.t(), gt_i_c), torch.matmul(pred_i_c.t(), pred_i_c)).solution
    ints_ratio1 = (gt_i_c - scale * pred_i_c).abs() / (gt_i_c + 1e-8)
    # Green channel:
    gt_i_c = gt_i[:, 1:2]
    pred_i_c = pred_i[:, 1:2]
    scale = torch.solve(torch.matmul(pred_i_c.t(), gt_i_c), torch.matmul(pred_i_c.t(), pred_i_c)).solution
    ints_ratio2 = (gt_i_c - scale * pred_i_c).abs() / (gt_i_c + 1e-8)
    # Blue channel:
    gt_i_c = gt_i[:, 2:3]
    pred_i_c = pred_i[:, 2:3]
    scale = torch.solve(torch.matmul(pred_i_c.t(), gt_i_c), torch.matmul(pred_i_c.t(), pred_i_c)).solution
    ints_ratio3 = (gt_i_c - scale * pred_i_c).abs() / (gt_i_c + 1e-8)

    ints_ratio = (ints_ratio1 + ints_ratio2 + ints_ratio3) / 3
    return ints_ratio.mean().item(), ints_ratio.mean(dim=-1)
//  [python]
def cal_ints_acc(gt_i, pred_i):
    # Red channel:
    gt_i_c = gt_i[:, :1]
    pred_i_c = pred_i[:, :1]
    scale = torch.linalg.lstsq(pred_i_c, gt_i_c).solution
    ints_ratio1 = (gt_i_c - scale * pred_i_c).abs() / (gt_i_c + 1e-8)
    # Green channel:
    gt_i_c = gt_i[:, 1:2]
    pred_i_c = pred_i[:, 1:2]
    scale = torch.linalg.lstsq(pred_i_c, gt_i_c).solution
    ints_ratio2 = (gt_i_c - scale * pred_i_c).abs() / (gt_i_c + 1e-8)
    # Blue channel:
    gt_i_c = gt_i[:, 2:3]
    pred_i_c = pred_i[:, 2:3]
    scale = torch.linalg.lstsq(pred_i_c, gt_i_c).solution
    ints_ratio3 = (gt_i_c - scale * pred_i_c).abs() / (gt_i_c + 1e-8)

    ints_ratio = (ints_ratio1 + ints_ratio2 + ints_ratio3) / 3
    return ints_ratio.mean().item(), ints_ratio.mean(dim=-1)
//  [python]
def Kron_reduction(edge_index,restriction_maps,boundary_vertices,interior_vertices):
    L = Laplacian(edge_index,restriction_maps)
    dv = restriction_maps.shape[-1]
    nbatch = L.shape[0]

    Bidx = get_matrix_indices(boundary_vertices,dv)
    Uidx = get_matrix_indices(interior_vertices,dv)
    batch_idx = np.arange(nbatch)

    LBB = L[np.ix_(batch_idx,Bidx,Bidx)]
    LUB = L[np.ix_(batch_idx,Uidx,Bidx)]
    LUU = L[np.ix_(batch_idx,Uidx,Uidx)]

    # lstsq recommended by pytorch docs https://pytorch.org/docs/stable/generated/torch.linalg.pinv.html
    schur = LBB - torch.transpose(LUB, 1, 2) @ torch.solve(LUB,LUU)[0]

    return schur
//  [python]
def Kron_reduction(edge_index,restriction_maps,boundary_vertices,interior_vertices):
    L = Laplacian(edge_index,restriction_maps)
    dv = restriction_maps.shape[-1]
    nbatch = L.shape[0]

    Bidx = get_matrix_indices(boundary_vertices,dv)
    Uidx = get_matrix_indices(interior_vertices,dv)
    batch_idx = np.arange(nbatch)

    LBB = L[np.ix_(batch_idx,Bidx,Bidx)]
    LUB = L[np.ix_(batch_idx,Uidx,Bidx)]
    LUU = L[np.ix_(batch_idx,Uidx,Uidx)]

    # lstsq recommended by pytorch docs https://pytorch.org/docs/stable/generated/torch.linalg.pinv.html
    schur = LBB - torch.transpose(LUB, 1, 2) @ torch.linalg.lstsq(LUU, LUB).solution

    return schur