jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_matrix_rank_atol(self, device, dtype):

        def run_test_atol(shape0, shape1, batch):
            a = make_tensor((*batch, shape0, shape1), dtype=dtype, device=device)
            # Check against NumPy output
            # Test float tol, and specific value for each matrix
            tolerances = [float(torch.rand(1)), ]
            # Test different types of tol tensor
            for tol_type in all_types():
                tolerances.append(make_tensor(a.shape[:-2], dtype=tol_type, device=device, low=0))
            # Test broadcasting of tol
            if a.ndim > 2:
                tolerances.append(make_tensor(a.shape[-3], dtype=torch.float32, device=device, low=0))
            for tol in tolerances:
                with pytorch_op_timer():
                    actual = torch.linalg.matrix_rank(a, atol=tol)
                with pytorch_op_timer():
                    actual_tol = torch.linalg.matrix_rank(a, tol=tol)
                self.assertEqual(actual, actual_tol)
                numpy_tol = tol if isinstance(tol, float) else tol.cpu().numpy()
                expected = np.linalg.matrix_rank(a.cpu().numpy(), tol=numpy_tol)
                self.assertEqual(actual, expected)

        shapes = (3, 13)
        batches = ((), (0, ), (4, ), (3, 5, ))
        for (shape0, shape1), batch in zip(itertools.product(shapes, reversed(shapes)), batches):
            run_test_atol(shape0, shape1, batch)
//  [python]
def test_matrix_rank_atol_rtol(self, device, dtype):
        make_fullrank = make_fullrank_matrices_with_distinct_singular_values
        make_arg = partial(make_fullrank, device=device, dtype=dtype)

        # creates a matrix with singular values rank=n and singular values in range [2/3, 3/2]
        # the singular values are 1 + 1/2, 1 - 1/3, 1 + 1/4, 1 - 1/5, ...
        n = 9
        a = make_arg(n, n)

        # test float and tensor variants
        for tol_value in [0.81, torch.tensor(0.81, device=device)]:
            # using rtol (relative tolerance) takes into account the largest singular value (1.5 in this case)
            with pytorch_op_timer():
                result = torch.linalg.matrix_rank(a, rtol=tol_value)
            self.assertEqual(result, 2)  # there are 2 singular values above 1.5*0.81 = 1.215

            # atol is used directly to compare with singular values
            with pytorch_op_timer():
                result = torch.linalg.matrix_rank(a, atol=tol_value)
            self.assertEqual(result, 7)  # there are 7 singular values above 0.81

            # when both are specified the maximum tolerance is used
            with pytorch_op_timer():
                result = torch.linalg.matrix_rank(a, atol=tol_value, rtol=tol_value)
            self.assertEqual(result, 2)
//  [python]
def solve_output(self, Hidden_States, y):
        t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
            self.device), Hidden_States), dim=1)
        HTH = t_hs.T @ t_hs
        HTY = t_hs.T @ y
        # ridge regression with pytorch
        I = (0 * torch.eye(HTH.size(0))).to(self.device)
        A = HTH + I
        # A = A.to(self.device)
        orig_rank = torch.matrix_rank(HTH).item()
        tag = 'Inverse' if orig_rank == t_hs.size(1) else 'Pseudo-inverse'
        if tag == 'Inverse':
            W = torch.mm(torch.inverse(A), HTY).t()
        else:
            W = torch.mm(torch.pinverse(A.cpu()),
                         HTY.cpu()).t().to(self.device)
        return W, tag
//  [python]
def solve_output(self, Hidden_States, y):
        t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
            self.device), Hidden_States), dim=1)
        HTH = t_hs.T @ t_hs
        HTY = t_hs.T @ y
        # ridge regression with pytorch
        I = (0 * torch.eye(HTH.size(0))).to(self.device)
        A = HTH + I
        # A = A.to(self.device)
        # orig_rank = torch.matrix_rank(HTH).item() # torch.matrix_rank is deprecated in favor of torch.linalg.matrix_rank in v1.9.0
        orig_rank = torch.linalg.matrix_rank(A.cpu(), hermitian=True).item()
        tag = 'Inverse' if orig_rank == t_hs.size(1) else 'Pseudo-inverse'
        if tag == 'Inverse':
            W = torch.mm(torch.inverse(A), HTY).t()
        else:
            W = torch.mm(torch.linalg.pinv(A.cpu()),
                         HTY.cpu()).t().to(self.device)
        return W, tag
//  [python]
def _generate_prob():
    m = 100
    n = 500

    A = torch.randn(m, n)
    A[0].abs_()
    rank = torch.matrix_rank(A)
    assert rank == m

    p = torch.randn(n).abs()  # Make positive.
    b = A @ p
    c = torch.randn(n)
    in_domain = lambda soln: torch.all(soln.x > 0)

    x = torch.randn(n).exp()  # Make positive.
    nu = torch.zeros(m)

    return Soln(x=x, nu=nu), LPCenteringProb(A=A, b=b, c=c, in_domain=in_domain)
//  [python]
def _generate_prob():
    m = 100
    n = 500

    A = torch.randn(m, n)
    A[0].abs_()
    rank = torch.linalg.matrix_rank(A)
    assert rank == m

    p = torch.randn(n).abs()  # Make positive.
    b = A @ p
    c = torch.randn(n)
    in_domain = lambda soln: torch.all(soln.x > 0)

    x = torch.randn(n).exp()  # Make positive.
    nu = torch.zeros(m)

    return Soln(x=x, nu=nu), LPCenteringProb(A=A, b=b, c=c, in_domain=in_domain)
//  [python]
def _generate_prob(infeasible=False, unbounded=False):
    m = 100
    n = 600

    if infeasible:
        A = torch.zeros(m, n)
        b = torch.ones(m)

        c = torch.randn(n)

        x = torch.randn(n).exp() * 3  # Make positive.
        nu = torch.randn(m)

    elif unbounded:
        A = torch.tensor([[1, 0, 0, 0], [0, 1, 0, 0]], dtype=torch.get_default_dtype())
        b = torch.tensor([0, 0], dtype=torch.get_default_dtype())
        c = torch.zeros(4)

        x = torch.randn(4).exp()
        nu = torch.zeros(2)
    else:
        A = torch.randn(m, n)
        A[0].abs_()
        rank = torch.linalg.matrix_rank(A)
        assert rank == m

        p = torch.randn(n).abs()  # Make positive.
        b = A @ p

        c = torch.randn(n)

        x = torch.randn(n).exp() * 3  # Make positive.
        nu = torch.randn(m)

    in_domain = lambda soln: torch.all(soln.x > 0)
    return Soln(x=x, nu=nu), LPCenteringProb(A=A, b=b, c=c, in_domain=in_domain)
//  [python]
def lstsq(A, B, require_grad):
    if require_grad:
        return torch.pinverse(A) @ B
    else:
        if A.dim() == 2:
            return torch.lstsq(B, A).solution[:A.size(-1)]
        elif A.dim() == 3:
            sols = [
                torch.lstsq(B[bidx], A[bidx]).solution[:A.size(-1)] for bidx in range(A.size(0))
            ]
            return torch.stack(sols, dim=0)
        else:
            raise NotImplementedError
//  [python]
def lstsq(A, B, require_grad):
    if require_grad:
        return torch.pinverse(A) @ B
    else:
        if A.dim() == 2:
            return torch.linalg.lstsq(A, B).solution[:A.size(-1)]
        elif A.dim() == 3:
            sols = [
                torch.linalg.lstsq(A[bidx], B[bidx]).solution[:A.size(-1)] for bidx in range(A.size(0))
            ]
            return torch.stack(sols, dim=0)
        else:
            raise NotImplementedError
//  [python]
def close_ho(device, Hidden_States, y, reg_lambda=0., report=False):
    
    map_size = Hidden_States.size(1)
    H = y.size(1)
    
    t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
        device), Hidden_States), dim=1)
    HTH = t_hs.T @ t_hs
    HTY = t_hs.T @ y
    # ridge regression with pytorch
    I = (reg_lambda * torch.eye(HTH.size(0))).to(device)
    A = HTH + I
    orig_rank = torch.matrix_rank(HTH).item()
    tag = 'Inverse'
    try:
        orig_rank = torch.matrix_rank(A).item()
        if orig_rank != t_hs.size(1):
            tag = 'Pseudo-inverse'
    except:
        tag = 'Pseudo-inverse'
    
    try:
        W = torch.linalg.lstsq(A.cpu(), HTY.cpu(), driver = 'gelsd').solution.T.to(device) # Attention this method may encounter bugs when above exception happens
    except:
        # If the computation of least square method meets bug, then make this subnetwork output zeros, tranferring the current error-feedback into the next construction iteration.
        W = torch.zeros(H, map_size+1).to(device)
        # When the training goes on, sometimes the input A may be singular, which makes the bug. This bug may be solved in pytorch 1.10. 

        # W = torch.mm(torch.linalg.pinv(A.cpu(), hermitian = True),
        #              HTY.cpu()).t().to(device)

    bias = nn.Parameter(W[:, 0], requires_grad=False)
    weight = nn.Parameter(W[:, 1:], requires_grad=False)
    if report:
        print('Global LSM: {} \t L2 regular: {}'.format(
            tag, 'True' if reg_lambda != 0 else 'False'))
    
    ho = fcLayer(map_size, H, device)
    ho.update(weight, bias)
    ho.freeze()
    ho.inverse = tag
    ho.reg_lambda = reg_lambda
    
    return ho
//  [python]
def close_ho(device, Hidden_States, y, reg_lambda=0., report=False):
    
    map_size = Hidden_States.size(1)
    H = y.size(1)
    
    t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
        device), Hidden_States), dim=1)
    HTH = t_hs.T @ t_hs
    HTY = t_hs.T @ y
    # ridge regression with pytorch
    I = (reg_lambda * torch.eye(HTH.size(0))).to(device)
    A = HTH + I
    # orig_rank = torch.matrix_rank(HTH).item() # torch.matrix_rank is deprecated in favor of torch.linalg.matrix_rank in v1.9.0
    tag = 'Inverse'
    try:
        orig_rank = torch.linalg.matrix_rank(A, hermitian=True).item()
        if orig_rank != t_hs.size(1):
            tag = 'Pseudo-inverse'
    except:
        tag = 'Pseudo-inverse'
    
    try:
        W = torch.linalg.lstsq(A.cpu(), HTY.cpu(), driver = 'gelsd').solution.T.to(device) # Attention this method may encounter bugs when above exception happens
    except:
        # If the computation of least square method meets bug, then make this subnetwork output zeros, tranferring the current error-feedback into the next construction iteration.
        W = torch.zeros(H, map_size+1).to(device)
        # When the training goes on, sometimes the input A may be singular, which makes the bug. This bug may be solved in pytorch 1.10. 

        # W = torch.mm(torch.linalg.pinv(A.cpu(), hermitian = True),
        #              HTY.cpu()).t().to(device)

    bias = nn.Parameter(W[:, 0], requires_grad=False)
    weight = nn.Parameter(W[:, 1:], requires_grad=False)
    if report:
        print('Global LSM: {} \t L2 regular: {}'.format(
            tag, 'True' if reg_lambda != 0 else 'False'))
    
    ho = fcLayer(map_size, H, device)
    ho.update(weight, bias)
    ho.freeze()
    ho.inverse = tag
    ho.reg_lambda = reg_lambda
    
    return ho
//  [python]
def __init__(
        self,
        d: int,
        t_span: Union[Tuple[float, float], Tuple[Tensor, Tensor]],
        n_reparam_samples: int,
        G: Tensor,
        drift: Callable[[Tensor, Tensor], Tensor],
        num_meas: int,
        measurement_noise: Tensor,
        tau: float = 1e-5,
        train_t: Tensor = None,
        train_x: Tensor = None,
        n_quad: int = 128,
        quad_percent: float = 0.8,
        n_tau: int = 200,
        Q_diag: Tensor = None,
    ) -> None:
        assert quad_percent > 0 and quad_percent < 1, "quad_percent must be in (0, 1)"
        if Q_diag is None:
            Q_diag = torch.ones(d) * 1e-0
            diffusion_prior = SparseDiagonalDiffusionPrior(
                d, Q_diag, n_reparam_samples, tau
            )
        else:
            diffusion_prior = ConstantDiagonalDiffusionPrior(d, Q_diag)
        fast_initialization = (
            (torch.matrix_rank(G) == d)
            and (train_t is not None)
            and (train_x is not None)
        )
        # transform = IdentityScaleTransform(d)
        if fast_initialization:
            stdev_transform = (torch.eye(d) - G).sum() == 0 and (train_x is not None)
            if stdev_transform:
                # transform = StdevScaleTransform(train_x)
                measurement_noise = measurement_noise  # * (1 / transform.scale ** 2)
            # train_x = transform(train_x)
            z = solve_least_squares(G, train_x.t(), gamma=1e-2).t()
            # , clamp_min=1e-1).t()
        else:
            train_t = None
            z = None
        marginal_sde = SpectralMarginalSDE(
            d,
            t_span,
            diffusion_prior=diffusion_prior,
            model_form="GLM",
            n_tau=n_tau,
            learn_inducing_locations=False,
            train_x=train_t,
            train_y=z,
        )
        quad_scheme = UnbiasedGaussLegendreQuad(
            t_span[0], t_span[1], n_quad, quad_percent=quad_percent
        )
        likelihood = IndepGaussLikelihood(G, num_meas, measurement_noise)
        if fast_initialization:
            m, dmdt = marginal_sde.mean(train_t, return_grad=True)
        else:
            m, dmdt = (None, None)
        sde_prior = ExactMotionModel(drift)
        super().__init__(
            marginal_sde,
            likelihood,
            quad_scheme,
            sde_prior,
            diffusion_prior,
            n_reparam_samples,
        )
//  [python]
def __init__(
        self,
        d: int,
        t_span: Union[Tuple[float, float], Tuple[Tensor, Tensor]],
        n_reparam_samples: int,
        G: Tensor,
        drift: Callable[[Tensor, Tensor], Tensor],
        num_meas: int,
        measurement_noise: Tensor,
        tau: float = 1e-5,
        train_t: Tensor = None,
        train_x: Tensor = None,
        n_quad: int = 128,
        quad_percent: float = 0.8,
        n_tau: int = 200,
        Q_diag: Tensor = None,
    ) -> None:
        assert quad_percent > 0 and quad_percent < 1, "quad_percent must be in (0, 1)"
        if Q_diag is None:
            Q_diag = torch.ones(d) * 1e-0
            diffusion_prior = SparseDiagonalDiffusionPrior(
                d, Q_diag, n_reparam_samples, tau
            )
        else:
            diffusion_prior = ConstantDiagonalDiffusionPrior(d, Q_diag)
        fast_initialization = (
            (torch.linalg.matrix_rank(G) == d)
            and (train_t is not None)
            and (train_x is not None)
        )
        # transform = IdentityScaleTransform(d)
        if fast_initialization:
            stdev_transform = (torch.eye(d) - G).sum() == 0 and (train_x is not None)
            if stdev_transform:
                # transform = StdevScaleTransform(train_x)
                measurement_noise = measurement_noise  # * (1 / transform.scale ** 2)
            # train_x = transform(train_x)
            z = solve_least_squares(G, train_x.t(), gamma=1e-2).t()
            # , clamp_min=1e-1).t()
        else:
            train_t = None
            z = None
        marginal_sde = SpectralMarginalSDE(
            d,
            t_span,
            diffusion_prior=diffusion_prior,
            model_form="GLM",
            n_tau=n_tau,
            learn_inducing_locations=False,
            train_x=train_t,
            train_y=z,
        )
        quad_scheme = UnbiasedGaussLegendreQuad(
            t_span[0], t_span[1], n_quad, quad_percent=quad_percent
        )
        likelihood = IndepGaussLikelihood(G, num_meas, measurement_noise)
        if fast_initialization:
            m, dmdt = marginal_sde.mean(train_t, return_grad=True)
        else:
            m, dmdt = (None, None)
        sde_prior = ExactMotionModel(drift)
        super().__init__(
            marginal_sde,
            likelihood,
            quad_scheme,
            sde_prior,
            diffusion_prior,
            n_reparam_samples,
        )
//  [python]
def __init__(
        self,
        d: int,
        t_span: Union[Tuple[float, float], Tuple[Tensor, Tensor]],
        degree: int,
        n_reparam_samples: int,
        G: Tensor,
        num_meas: int,
        measurement_noise: Tensor,
        tau: float = 1e-5,
        train_t: Tensor = None,
        train_x: Tensor = None,
        input_labels: List[str] = None,
        n_quad: int = 128,
        quad_percent: float = 0.8,
        n_tau: int = 200,
    ) -> None:
        assert quad_percent > 0 and quad_percent < 1, "quad_percent must be in (0, 1)"
        Q_diag = torch.ones(d) * 1e-0
        diffusion_prior = SparseDiagonalDiffusionPrior(
            d, Q_diag, n_reparam_samples, tau
        )
        fast_initialization = (
            (torch.linalg.matrix_rank(G) == d)
            and (train_t is not None)
            and (train_x is not None)
        )
        # transform = IdentityScaleTransform(d)
        if fast_initialization:
            stdev_transform = (torch.eye(d) - G).sum() == 0 and (train_x is not None)
            if stdev_transform:
                # transform = StdevScaleTransform(train_x)
                measurement_noise = measurement_noise  # * (1 / transform.scale ** 2)
            # train_x = transform(train_x)
            z = solve_least_squares(G, train_x.t(), gamma=1e-2).t()
            # , clamp_min=1e-1).t()
        else:
            train_t = None
            z = None
        marginal_sde = SpectralMarginalSDE(
            d,
            t_span,
            diffusion_prior=diffusion_prior,
            model_form="GLM",
            n_tau=n_tau,
            learn_inducing_locations=False,
            train_x=train_t,
            train_y=z,
        )
        quad_scheme = UnbiasedGaussLegendreQuad(
            t_span[0], t_span[1], n_quad, quad_percent=quad_percent
        )
        likelihood = IndepGaussLikelihood(G, num_meas, measurement_noise)
        features = SparsePolynomialFeatures(d, degree=degree, input_labels=input_labels)
        if fast_initialization:
            m, dmdt = marginal_sde.mean(train_t, return_grad=True)
        else:
            m, dmdt = (None, None)
        sde_prior = SparseMultioutputGLM(
            d,
            SparseFeatures=features,
            n_reparam_samples=n_reparam_samples,
            tau=tau,
            train_x=m,
            train_y=dmdt,
            # transform=transform,
        )
        super().__init__(
            marginal_sde,
            likelihood,
            quad_scheme,
            sde_prior,
            diffusion_prior,
            n_reparam_samples,
        )
//  [python]
def forward(self, img):
    """Forward function of SIIE net."""
    if self._initial_est:
      initial_ill = torch.mean(img, dim=[2, 3])
      gains = torch.unsqueeze(initial_ill[:, 1], dim=-1) / initial_ill
      img[:, 0, ...] *= torch.unsqueeze(torch.unsqueeze(gains[:, 0], axis=-1),
                                        axis=-1)
      img[:, 2, ...] *= torch.unsqueeze(torch.unsqueeze(gains[:, 2], axis=-1),
                                        axis=-1)

    m = self._sensor_mapping_net(img)
    m = m.reshape([-1, 3, 3])
    m = torch.abs(m) / (
        torch.unsqueeze(torch.sum(torch.abs(m), axis=-1), axis=-1) + EPS)
    img_mapped = img.clone()
    for i in range(m.shape[0]):
      temp = torch.mm(torch.squeeze(m[i, :, :]),
                      torch.reshape(torch.squeeze(
                          img[i, :, :, :]), (3, -1)))
      img_mapped[i, :, :, :] = torch.reshape(temp, (img.shape[1], img.shape[2],
                                                    img.shape[3]))

    ill_mapped = self._illum_est_net(img_mapped)
    inds = (torch.matrix_rank(m) < 3).to(device=m.device)
    offset = torch.zeros(m.shape[0], 1).to(device=m.device)
    if torch.sum(inds) > 0:
      offset[inds, ...] = torch.unsqueeze(
          torch.rand(torch.sum(inds)), axis=-1).to(device=m.device) * 0.0001
    inv_m = torch.linalg.inv(m + torch.unsqueeze(offset, axis=-1))
    ill = ill_mapped.clone()
    for i in range(inv_m.shape[0]):
      temp = torch.mm(torch.squeeze(inv_m[i, :, :]),
                      torch.unsqueeze(ill_mapped[i, :], axis=-1))
      ill[i, :] = torch.reshape(temp, ill[i, :].shape)
    if self._initial_est:
      ill = initial_ill * ill
    return ill, m
//  [python]
def forward(self, img):
    """Forward function of SIIE net."""
    if self._initial_est:
      initial_ill = torch.mean(img, dim=[2, 3])
      gains = torch.unsqueeze(initial_ill[:, 1], dim=-1) / initial_ill
      img[:, 0, ...] *= torch.unsqueeze(torch.unsqueeze(gains[:, 0], axis=-1),
                                        axis=-1)
      img[:, 2, ...] *= torch.unsqueeze(torch.unsqueeze(gains[:, 2], axis=-1),
                                        axis=-1)

    m = self._sensor_mapping_net(img)
    m = m.reshape([-1, 3, 3])
    m = torch.abs(m) / (
        torch.unsqueeze(torch.sum(torch.abs(m), axis=-1), axis=-1) + EPS)
    img_mapped = img.clone()
    for i in range(m.shape[0]):
      temp = torch.mm(torch.squeeze(m[i, :, :]),
                      torch.reshape(torch.squeeze(
                          img[i, :, :, :]), (3, -1)))
      img_mapped[i, :, :, :] = torch.reshape(temp, (img.shape[1], img.shape[2],
                                                    img.shape[3]))

    ill_mapped = self._illum_est_net(img_mapped)
    inds = (torch.linalg.matrix_rank(m) < 3).to(device=m.device)
    offset = torch.zeros(m.shape[0], 1).to(device=m.device)
    if torch.sum(inds) > 0:
      offset[inds, ...] = torch.unsqueeze(
          torch.rand(torch.sum(inds)), axis=-1).to(device=m.device) * 0.0001
    inv_m = torch.linalg.inv(m + torch.unsqueeze(offset, axis=-1))
    ill = ill_mapped.clone()
    for i in range(inv_m.shape[0]):
      temp = torch.mm(torch.squeeze(inv_m[i, :, :]),
                      torch.unsqueeze(ill_mapped[i, :], axis=-1))
      ill[i, :] = torch.reshape(temp, ill[i, :].shape)
    if self._initial_est:
      ill = initial_ill * ill
    return ill, m
//  [python]
def __init__(
        self,
        d: int,
        t_span: Union[Tuple[float, float], Tuple[Tensor, Tensor]],
        n_reparam_samples: int,
        G: Tensor,
        drift_layer_description: List[int],
        nonlinearity: nn.Module,
        measurement_noise: Tensor,
        tau: float = 1e-5,
        train_t: Optional[Tensor] = None,
        train_x: Optional[Tensor] = None,
        n_quad: int = 128,
        quad_percent: float = 0.8,
        n_tau: int = 200,
        Q_diag: Optional[Tensor] = None,
    ) -> None:
        assert quad_percent > 0 and quad_percent < 1, "quad_percent must be in (0, 1)"
        if Q_diag is None:
            Q_diag = torch.ones(d) * 1e-0
            diffusion_prior = SparseDiagonalDiffusionPrior(
                d, Q_diag, n_reparam_samples, tau
            )
        else:
            diffusion_prior = ConstantDiagonalDiffusionPrior(d, Q_diag)
        fast_initialization = (
            (torch.linalg.matrix_rank(G) == d)
            and (train_t is not None)
            and (train_x is not None)
        )
        if fast_initialization:
            assert isinstance(train_x, Tensor)  # helping the type checker
            z = solve_least_squares(G, train_x.t(), gamma=1e-2).t()
        else:
            train_t = None
            z = None
        marginal_sde = DiagonalMarginalSDE(
            d,
            t_span,
            diffusion_prior=diffusion_prior,
            model_form="GLM",
            n_tau=n_tau,
            learn_inducing_locations=False,
            train_x=train_t,
            train_y=z,
        )
        quad_scheme = UnbiasedGaussLegendreQuad(
            t_span[0], t_span[1], n_quad, quad_percent=quad_percent
        )
        likelihood = IndepGaussLikelihood(G, G.shape[0], measurement_noise)
        sde_prior = DriftFCNN(d, drift_layer_description, nonlinearity)
        super().__init__(
            marginal_sde,
            likelihood,
            quad_scheme,
            sde_prior,
            diffusion_prior,
            n_reparam_samples,
        )
//  [python]
def close_ho(device, Hidden_States, y, lambda_reg=0., report=False):
    
    map_size = Hidden_States.size(1)
    H = y.size(1)
    
    t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
        device), Hidden_States), dim=1)
    HTH = t_hs.T @ t_hs
    HTY = t_hs.T @ y
    # ridge regression with pytorch
    I = (lambda_reg * torch.eye(HTH.size(0))).to(device)
    A = HTH + I
    orig_rank = torch.matrix_rank(HTH).item()
    tag = 'Inverse'
    try:
        orig_rank = torch.matrix_rank(A).item()
        if orig_rank != t_hs.size(1):
            tag = 'Pseudo-inverse'
    except:
        tag = 'Pseudo-inverse'
    
    try:
        W = torch.linalg.lstsq(A.cpu(), HTY.cpu(), driver = 'gelsd').solution.T.to(device) # Attention this method may encounter bugs when above exception happens
    except:
        # If the computation of least square method meets bug, then make this subnetwork output zeros, tranferring the current error-feedback into the next construction iteration.
        W = torch.zeros(H, map_size+1).to(device)
        # When the training goes on, sometimes the input A may be singular, which makes the bug. This bug may be solved in pytorch 1.10. 

        # W = torch.mm(torch.linalg.pinv(A.cpu(), hermitian = True),
        #              HTY.cpu()).t().to(device)

    bias = nn.Parameter(W[:, 0], requires_grad=False)
    weight = nn.Parameter(W[:, 1:], requires_grad=False)
    if report:
        print('Global LSM: {} \t L2 regular: {}'.format(
            tag, 'True' if lambda_reg != 0 else 'False'))
    
    ho = fcLayer(map_size, H, device)
    ho.update(weight, bias)
    ho.freeze()
    ho.inverse = tag
    ho.lambda_reg = lambda_reg
    
    return ho
//  [python]
def close_ho(device, Hidden_States, y, lambda_reg=0., report=False):
    
    map_size = Hidden_States.size(1)
    H = y.size(1)
    
    t_hs = torch.cat((torch.ones(Hidden_States.size(0), 1).to(
        device), Hidden_States), dim=1)
    HTH = t_hs.T @ t_hs
    HTY = t_hs.T @ y
    # ridge regression with pytorch
    I = (lambda_reg * torch.eye(HTH.size(0))).to(device)
    A = HTH + I
    # orig_rank = torch.matrix_rank(HTH).item() # torch.matrix_rank is deprecated in favor of torch.linalg.matrix_rank in v1.9.0
    tag = 'Inverse'
    try:
        orig_rank = torch.linalg.matrix_rank(A, hermitian=True).item()
        if orig_rank != t_hs.size(1):
            tag = 'Pseudo-inverse'
    except:
        tag = 'Pseudo-inverse'
    
    try:
        W = torch.linalg.lstsq(A.cpu(), HTY.cpu(), driver = 'gelsd').solution.T.to(device) # Attention this method may encounter bugs when above exception happens
    except:
        # If the computation of least square method meets bug, then make this subnetwork output zeros, tranferring the current error-feedback into the next construction iteration.
        W = torch.zeros(H, map_size+1).to(device)
        # When the training goes on, sometimes the input A may be singular, which makes the bug. This bug may be solved in pytorch 1.10. 

        # W = torch.mm(torch.linalg.pinv(A.cpu(), hermitian = True),
        #              HTY.cpu()).t().to(device)

    bias = nn.Parameter(W[:, 0], requires_grad=False)
    weight = nn.Parameter(W[:, 1:], requires_grad=False)
    if report:
        print('Global LSM: {} \t L2 regular: {}'.format(
            tag, 'True' if lambda_reg != 0 else 'False'))
    
    ho = fcLayer(map_size, H, device)
    ho.update(weight, bias)
    ho.freeze()
    ho.inverse = tag
    ho.lambda_reg = lambda_reg
    
    return ho
//  [python]
def calIntsAcc(gt_i, pred_i, data_batch=1):
    # import pdb
    # pdb.set_trace()
    n, c, h, w = gt_i.shape
    pred_i  = pred_i.view(n, c, h, w)
    ref_int = gt_i[:, :3].repeat(1, gt_i.shape[1] // 3, 1, 1)
    gt_i  = gt_i / ref_int
    scale = torch.lstsq(gt_i.view(-1, 1), pred_i.view(-1, 1))
    ints_ratio = (gt_i - scale[0][0] * pred_i).abs() / (gt_i + 1e-8)
    ints_error = torch.stack(ints_ratio.split(3, 1), 1).mean(2)
    return {'ints_ratio': ints_ratio.mean().item()}, ints_error.squeeze()
//  [python]
def calIntsAcc(gt_i, pred_i, data_batch=1):
    n, c, h, w = gt_i.shape
    pred_i  = pred_i.view(n, c, h, w)
    ref_int = gt_i[:, :3].repeat(1, gt_i.shape[1] // 3, 1, 1)
    gt_i  = gt_i / ref_int
    scale = torch.linalg.lstsq(gt_i.view(-1, 1), pred_i.view(-1, 1))
    ints_ratio = (gt_i - scale[0][0] * pred_i).abs() / (gt_i + 1e-8)
    ints_error = torch.stack(ints_ratio.split(3, 1), 1).mean(2)
    return {'ints_ratio': ints_ratio.mean().item()}, ints_error.squeeze()
//  [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)

        # Remove blank (all-zero) episodes that only exist because episode lengths vary
        flat_mask = episodes.mask.flatten()
        flat_mask_nnz = torch.nonzero(flat_mask)
        featmat = featmat[flat_mask_nnz].view(-1, self.feature_size)
        returns = returns[flat_mask_nnz].view(-1, 1)

        reg_coeff = self._reg_coeff
        XT_y = torch.matmul(featmat.t(), returns)
        XT_X = torch.matmul(featmat.t(), featmat)
        for _ in range(5):
            try:
                coeffs, _ = torch.lstsq(XT_y, XT_X + reg_coeff * self._eye)

                # An extra round of increasing regularization eliminated
                # inf or nan in the least-squares solution most of the time
                if torch.isnan(coeffs).any() or torch.isinf(coeffs).any():
                    raise RuntimeError

                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.weight.copy_(coeffs.flatten())
//  [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)

        # Remove blank (all-zero) episodes that only exist because episode lengths vary
        flat_mask = episodes.mask.flatten()
        flat_mask_nnz = torch.nonzero(flat_mask)
        featmat = featmat[flat_mask_nnz].view(-1, self.feature_size)
        returns = returns[flat_mask_nnz].view(-1, 1)

        reg_coeff = self._reg_coeff
        XT_y = torch.matmul(featmat.t(), returns)
        XT_X = torch.matmul(featmat.t(), featmat)
        for _ in range(5):
            try:
                coeffs, _ = torch.linalg.lstsq(XT_y, XT_X + reg_coeff * self._eye)

                # An extra round of increasing regularization eliminated
                # inf or nan in the least-squares solution most of the time
                if torch.isnan(coeffs).any() or torch.isinf(coeffs).any():
                    raise RuntimeError

                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.weight.copy_(coeffs.flatten())
//  [python]
def least_squares_qr(A, B):
    """Faster lstsq variant that uses PyTorch native QR factorization"""
    m, n = A.shape
    _, k = B.shape
    output, _ = torch.lstsq(B, A)
    X = output[:n]
    residuals = output[n:].square().sum(0)
    return X, residuals
//  [python]
def least_squares_qr(A, B):
    """Faster lstsq variant that uses PyTorch native QR factorization"""
    m, n = A.shape
    _, k = B.shape
    output, _ = torch.linalg.lstsq(B, A)
    X = output[:n]
    residuals = output[n:].square().sum(0)
    return X, residuals
//  [python]
def covariance_adjusted_query_key(self, k):
        if len(k.shape) == 1:
            return torch.lstsq(k[:, None], self.c_matrix)[0][:, 0]
        return torch.lstsq(k.permute(1, 0), self.c_matrix)[0].permute(1, 0)
//  [python]
def covariance_adjusted_query_key(self, k):
        if len(k.shape) == 1:
            return torch.linalg.lstsq(k[:, None], self.c_matrix)[0][:, 0]
        return torch.linalg.lstsq(k.permute(1, 0), self.c_matrix)[0].permute(1, 0)
//  [python]
def covariance_adjusted_query_key(self, k):
        # Returns C^{-1}k, computing it more stably than literally inverting C.
        if len(k.shape) == 1:
            return torch.lstsq(k[:, None], self.c_matrix)[0][:, 0]
        return torch.lstsq(k.permute(1, 0), self.c_matrix)[0].permute(1, 0)
//  [python]
def covariance_adjusted_query_key(self, k):
        # Returns C^{-1}k, computing it more stably than literally inverting C.
        if len(k.shape) == 1:
            return torch.linalg.lstsq(self.c_matrix, k[:, None])[0][:, 0]
        return torch.linalg.lstsq(self.c_matrix, k.permute(1, 0))[0].permute(1, 0)
//  [python]
def ProjectFunctionNeRFSparse(
    order: int,
    spherical_function: 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_function(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, device=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_function: 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_function(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, device=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 color_shift_remove(self, 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(self, 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 covariance_adjusted_key(self, k, kout):
        assert 'adain' in self.firstlayer
        # The idea here is that we should be constrained to the low
        # rank subspace corresponding to (CS)^{-1} k.
        assert kout.style.shape[0] == 1
        # cs = self.c_matrix * kout.style[0][:,None]  # SC version
        cs = self.c_matrix * kout.style[0][None, :]  # CS version, correct I think
        if len(k.shape) == 1:
            return torch.lstsq(k[:, None], cs)[0][:, 0]
        return torch.lstsq(k.permute(1, 0), cs)[0].permute(1, 0)
//  [python]
def covariance_adjusted_key(self, k, kout):
        assert 'adain' in self.firstlayer
        # The idea here is that we should be constrained to the low
        # rank subspace corresponding to (CS)^{-1} k.
        assert kout.style.shape[0] == 1
        # cs = self.c_matrix * kout.style[0][:,None]  # SC version
        cs = self.c_matrix * kout.style[0][None, :]  # CS version, correct I think
        if len(k.shape) == 1:
            return torch.linalg.lstsq(k[:, None], cs)[0][:, 0]
        return torch.linalg.lstsq(k.permute(1, 0), cs)[0].permute(1, 0)
//  [python]
def least_sq_with_known_values(A, b, known=None):
	"""Solves the least squares problem minx ||Ax - b||2, where some values of x are known.
	Works by moving all known variables from A to b.

	:param A: full rank matrix of size (mxn)
	:param b: matrix of size (m x k)
	:param known: dict of known_variable : value

	:type A: torch.Tensor
	:type B: torch.Tensor
	:type known: dict

	:returns x: matrix of size (n x k)
	"""

	M, N = A.shape
	M2, K = b.shape
	assert M == M2, "A's first dimension must match b's second"

	if known is None: known = {}

	# Move to b
	for index, val in known.items():
		col = A[:, index]
		b -= torch.einsum("i,j->ij", col, val)

	# Remove from A
	unknown = [n for n in range(N) if n not in known]
	A = A[:, unknown]  # only continue with cols for unknowns

	x, QR = torch.lstsq(b, A)

	# all unknown values have now been found. Now create the output tensor, of known and unknown values in correct positions
	x_out = torch.zeros((N, K))
	if known is not None:
		# Assign initially known values to x_out
		for index, val in known.items():
			x_out[index] = val

		# Assign initially unknown values to x_out
		x_out[unknown] = x[:len(unknown)]

	## X has shape max(m, n) x k. Only want first n rows
	return x_out
//  [python]
def least_sq_with_known_values(A, b, known=None):
	"""Solves the least squares problem minx ||Ax - b||2, where some values of x are known.
	Works by moving all known variables from A to b.

	:param A: full rank matrix of size (mxn)
	:param b: matrix of size (m x k)
	:param known: dict of known_variable : value

	:type A: torch.Tensor
	:type B: torch.Tensor
	:type known: dict

	:returns x: matrix of size (n x k)
	"""

	M, N = A.shape
	M2, K = b.shape
	assert M == M2, "A's first dimension must match b's second"

	if known is None: known = {}

	# Move to b
	for index, val in known.items():
		col = A[:, index]
		b -= torch.einsum("i,j->ij", col, val)

	# Remove from A
	unknown = [n for n in range(N) if n not in known]
	A = A[:, unknown]  # only continue with cols for unknowns

	x, QR = torch.linalg.lstsq(A, b)

	# all unknown values have now been found. Now create the output tensor, of known and unknown values in correct positions
	x_out = torch.zeros((N, K))
	if known is not None:
		# Assign initially known values to x_out
		for index, val in known.items():
			x_out[index] = val

		# Assign initially unknown values to x_out
		x_out[unknown] = x[:len(unknown)]

	## X has shape max(m, n) x k. Only want first n rows
	return x_out
//  [python]
def __find_concentration(self, OD, HE):
        # rows correspond to channels (RGB), columns to OD values
        Y = OD.T

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

        # determine concentrations of the individual stains
        return torch.linalg.lstsq(Y, HE)[0][:2]
//  [python]
def match_colors(im_ref, im_q, im_test):

    im_ref_mean_re = im_ref.view(*im_ref.shape[:2], -1)
    im_q_mean_re = im_q.view(*im_q.shape[:2], -1)

    # Estimate color transformation matrix by minimizing the least squares error
    c_mat_all = []
    for ir, iq in zip(im_ref_mean_re, im_q_mean_re):
        c = torch.lstsq(ir.t(), iq.t())
        c = c.solution[:im_ref_mean_re.size(1)]
        c_mat_all.append(c)

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

    return im_t_conv
//  [python]
def match_colors(im_ref, im_q, im_test):

    im_ref_mean_re = im_ref.view(*im_ref.shape[:2], -1)
    im_q_mean_re = im_q.view(*im_q.shape[:2], -1)

    # Estimate color transformation matrix by minimizing the least squares error
    c_mat_all = []
    for ir, iq in zip(im_ref_mean_re, im_q_mean_re):
        c = torch.linalg.lstsq(ir.t(), iq.t())
        c = c.solution[:im_ref_mean_re.size(1)]
        c_mat_all.append(c)

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

    return im_t_conv
//  [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 _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()
            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 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.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