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

            solve_g = torch.linalg.solve(g_reg, v)
            solve_a = torch.linalg.solve(a_reg, solve_g.t()).t()
            if layer.bias is None:
                solve_tuple = (solve_a.view(*sw),)
            else:
                solve_tuple = (solve_a[:, :-1].contiguous().view(*sw),
                               solve_a[:, -1].contiguous())
            out_dict[layer_id] = solve_tuple
        return PVector(layer_collection=vs.layer_collection,
                       dict_repr=out_dict)
//  [python]
def _inner(self, X, G, H):
        a = torch.solve(G, X)
        b = torch.solve(H, X)
        return torch.sum(a * b)
//  [python]
def _inner(self, X, G, H):
        a = torch.linalg.solve(X, G)
        b = torch.linalg.solve(X, H)
        return torch.sum(a * b)
//  [python]
def solve_wrapper(b, A):
    if A.numel() > 0:
        return torch.solve(b, A)[0]
    else:
        return b + A.reshape_as(b)
//  [python]
def solve_wrapper(b, A):
    if A.numel() > 0:
        return torch.linalg.solve(A, b)
    else:
        return b + A.reshape_as(b)
//  [python]
def forward(self, x):
        skip = x.clone()
        att = self.att(x)
        grad = self.grad(x)


        se = self.se(x)

        n, c, h, w = x.shape

        att = att.reshape(n*self.gr, 4, h*w, 1).permute(0, 2, 1, 3)
        grad = grad.reshape(n*self.gr, 4, h*w, 1).permute(0, 2, 1, 3)

        A = self.a * att
        B = grad * att

        A = A.reshape(n*self.gr, h*w*4, h*w)
        B = B.reshape(n*self.gr, h*w*4, 1)

        AT = A.permute(0, 2, 1)

        ATA = torch.bmm(AT, A)
        ATB = torch.bmm(AT, B)

        jitter = torch.eye(n=h*w, dtype=x.dtype, device=x.device).unsqueeze(0) * 1e-12
        x, _ = torch.solve(ATB, ATA+jitter)

        x = x.reshape(n, self.gr, h, w)

        x = self.ins(x)

        x = se * x

        x = self.post(x)

        return x
//  [python]
def forward(self, x):
        skip = x.clone()
        att = self.att(x)
        grad = self.grad(x)


        se = self.se(x)

        n, c, h, w = x.shape

        att = att.reshape(n*self.gr, 4, h*w, 1).permute(0, 2, 1, 3)
        grad = grad.reshape(n*self.gr, 4, h*w, 1).permute(0, 2, 1, 3)

        A = self.a * att
        B = grad * att

        A = A.reshape(n*self.gr, h*w*4, h*w)
        B = B.reshape(n*self.gr, h*w*4, 1)

        AT = A.permute(0, 2, 1)

        ATA = torch.bmm(AT, A)
        ATB = torch.bmm(AT, B)

        jitter = torch.eye(n=h*w, dtype=x.dtype, device=x.device).unsqueeze(0) * 1e-12
        x = torch.linalg.solve(ATA+jitter, ATB)

        x = x.reshape(n, self.gr, h, w)

        x = self.ins(x)

        x = se * x

        x = self.post(x)

        return x
//  [python]
def test_old_solve(self, device, dtype):
        for (k, n) in zip([2, 3, 5], [3, 5, 7]):
            b, A = self.solve_test_helper((n, n), (n, k), device, dtype)
            x = torch.solve(b, A)[0]
            self.assertEqual(b, np.matmul(A.cpu(), x.cpu()))
//  [python]
def test_old_solve(self, device, dtype):
        for (k, n) in zip([2, 3, 5], [3, 5, 7]):
            b, A = self.solve_test_helper((n, n), (n, k), device, dtype)
            x = torch.linalg.solve(A, b)
            self.assertEqual(b, np.matmul(A.cpu(), x.cpu()))
//  [python]
def M(self, q, eps=1e-4):
        """Returns a function that multiplies the mass matrix M(q) by a vector qdot

        Args:
            q: bs x D Tensor representing the position
        """
        assert q.ndim == 2
        lower_triangular = self.tril_Minv(q)
        assert lower_triangular.ndim == 3

        def M_func(qdot):
            assert qdot.ndim == 2
            qdot = qdot.unsqueeze(-1)
            diag_noise = eps * torch.eye(lower_triangular.size(-1), dtype=qdot.dtype, device=qdot.device)
            M_times_qdot = torch.solve(
                    qdot,
                    lower_triangular @ lower_triangular.transpose(-2, -1) + diag_noise
            ).solution.squeeze(-1)
            return M_times_qdot

        return M_func
//  [python]
def M(self, q, eps=1e-4):
        """Returns a function that multiplies the mass matrix M(q) by a vector qdot

        Args:
            q: bs x D Tensor representing the position
        """
        assert q.ndim == 2
        lower_triangular = self.tril_Minv(q)
        assert lower_triangular.ndim == 3

        def M_func(qdot):
            assert qdot.ndim == 2
            qdot = qdot.unsqueeze(-1)
            diag_noise = eps * torch.eye(lower_triangular.size(-1), dtype=qdot.dtype, device=qdot.device)
            M_times_qdot = torch.linalg.solve(
                    lower_triangular @ lower_triangular.transpose(-2, -1) + diag_noise,
                    qdot
            ).squeeze(-1)
            return M_times_qdot

        return M_func
//  [python]
def kalman_gain(self, cov: Tensor, H: Tensor, R: Tensor) -> Tensor:
        Ht = H.permute(0, 2, 1)
        system_covariance = H.matmul(cov).matmul(Ht) + R
        covs_measured = cov.matmul(Ht)
        A = system_covariance.permute(0, 2, 1)
        B = covs_measured.permute(0, 2, 1)
        Kt, _ = torch.solve(B, A)
        K = Kt.permute(0, 2, 1)
        return K
//  [python]
def kalman_gain(self, cov: Tensor, H: Tensor, R: Tensor) -> Tensor:
        Ht = H.permute(0, 2, 1)
        system_covariance = H.matmul(cov).matmul(Ht) + R
        covs_measured = cov.matmul(Ht)
        A = system_covariance.permute(0, 2, 1)
        B = covs_measured.permute(0, 2, 1)
        Kt = torch.linalg.solve(A, B)
        K = Kt.permute(0, 2, 1)
        return K
//  [python]
def forward(self, input, weight, logdet=None, reverse=False):
        """
        log-det = log|abs(|W|)| * pixels
        """
        assert input.shape[1] ** 2 == weight.shape[1], (input.shape, weight.shape)
        assert input.shape[2:] == weight.shape[2:] or input.shape[2:] == (1, 1)

        weight, weight_logdet = self.get_weight(weight)
        if not reverse:
            z = torch.solve(input.permute(0, 2, 3, 1).unsqueeze(-1), weight)[0].permute(0, 3, 1, 2, 4).squeeze(
                -1).contiguous()
            if logdet is not None:
                logdet = logdet - weight_logdet
            return z, logdet
        else:
            z = torch.matmul(weight, input.permute(0, 2, 3, 1).unsqueeze(-1)).permute(0, 3, 1, 2, 4).squeeze(
                -1).contiguous()
            if logdet is not None:
                logdet = logdet + weight_logdet
            return z, logdet
//  [python]
def forward(self, input, weight, logdet=None, reverse=False):
        """
        log-det = log|abs(|W|)| * pixels
        """
        assert input.shape[1] ** 2 == weight.shape[1], (input.shape, weight.shape)
        assert input.shape[2:] == weight.shape[2:] or input.shape[2:] == (1, 1)

        weight, weight_logdet = self.get_weight(weight)
        if not reverse:
            z = torch.linalg.solve(weight, input.permute(0, 2, 3, 1).unsqueeze(-1)).permute(0, 3, 1, 2, 4).squeeze(-1).contiguous()
            if logdet is not None:
                logdet = logdet - weight_logdet
            return z, logdet
        else:
            z = torch.matmul(weight, input.permute(0, 2, 3, 1).unsqueeze(-1)).permute(0, 3, 1, 2, 4).squeeze(-1).contiguous()
            if logdet is not None:
                logdet = logdet + weight_logdet
            return z, logdet
//  [python]
def test_old_solve_batched_many_batches(self, device, dtype):
        for A_dims, b_dims in zip([(256, 256, 5, 5), (3, 3)], [(5, 1), (512, 512, 3, 1)]):
            b, A = self.solve_test_helper(A_dims, b_dims, device, dtype)
            x, _ = torch.solve(b, A)
            Ax = torch.matmul(A, x)
            self.assertEqual(Ax, b.expand_as(x))
//  [python]
def test_old_solve_batched_many_batches(self, device, dtype):
        for A_dims, b_dims in zip([(256, 256, 5, 5), (3, 3)], [(5, 1), (512, 512, 3, 1)]):
            b, A = self.solve_test_helper(A_dims, b_dims, device, dtype)
            x = torch.linalg.solve(A, b)
            Ax = torch.matmul(A, x)
            self.assertEqual(Ax, b.expand_as(x))
//  [python]
def forward(self,x):
        """
        computes the PCMC choice probabilities P(.,S)

        Args:
        x- indicator vector for choice set S, i.e. a 'size(S)-hot' encoding of the
        choice set, or a batch of such encodings
        """
        #batching vectorizes in a peculiar way for PCMC, but works automatically
        #for other models, so store whether we plan to use batch at model initialization
        #and encode it here
        if self.batch:

            L = [] #empty list that will later contain choice probs

            for S in x.split(1):
                S=x
                S_mat = torch.mm(torch.t(S),S)
                Q = torch.clamp(self.Q*S_mat,min=self.epsilon)
                for i in range(self.n-1):
                    Q[i,i]=-torch.sum(Q[i,:])+Q[i,i]
                Q[:,self.n-1]=S
                b = Variable(torch.zeros(self.n)+self.epsilon)
                b[-1]=1
                pi,LU = torch.gesv(b,torch.t(Q))
                pi = S*torch.t(pi)+1e-16
                L.append(torch.log(pi/torch.sum(pi)))
            return torch.cat(L)
        else:
            S=x
            S_mat = torch.mm(torch.t(S),S)
            Q = torch.clamp(self.Q*S_mat,min=self.epsilon)
            for i in range(self.n-1):
                Q[i,i]=-torch.sum(Q[i,:])+Q[i,i]
            Q[:,self.n-1]=S
            b = Variable(torch.zeros(self.n,1))
            b[-1]=1
            pi,LU = torch.solve(b,torch.t(Q))
            pi = S*torch.t(pi)+1e-16

            return torch.log(pi/torch.sum(pi))
//  [python]
def forward(self,x):
        """
        computes the PCMC choice probabilities P(.,S)

        Args:
        x- indicator vector for choice set S, i.e. a 'size(S)-hot' encoding of the
        choice set, or a batch of such encodings
        """
        #batching vectorizes in a peculiar way for PCMC, but works automatically
        #for other models, so store whether we plan to use batch at model initialization
        #and encode it here
        if self.batch:

            L = [] #empty list that will later contain choice probs

            for S in x.split(1):
                S=x
                S_mat = torch.mm(torch.t(S),S)
                Q = torch.clamp(self.Q*S_mat,min=self.epsilon)
                for i in range(self.n-1):
                    Q[i,i]=-torch.sum(Q[i,:])+Q[i,i]
                Q[:,self.n-1]=S
                b = Variable(torch.zeros(self.n)+self.epsilon)
                b[-1]=1
                pi,LU = torch.linalg.solve(torch.t(Q),b)
                pi = S*torch.t(pi)+1e-16
                L.append(torch.log(pi/torch.sum(pi)))
            return torch.cat(L)
        else:
            S=x
            S_mat = torch.mm(torch.t(S),S)
            Q = torch.clamp(self.Q*S_mat,min=self.epsilon)
            for i in range(self.n-1):
                Q[i,i]=-torch.sum(Q[i,:])+Q[i,i]
            Q[:,self.n-1]=S
            b = Variable(torch.zeros(self.n,1))
            b[-1]=1
            pi,LU = torch.linalg.solve(torch.t(Q),b)
            pi = S*torch.t(pi)+1e-16

            return torch.log(pi/torch.sum(pi))
//  [python]
def test_old_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve

        def run_test(A_dims, b_dims):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            b, A = self.solve_test_helper(A_batch_dims + (A_matrix_size, A_matrix_size), b_dims, device, dtype)
            x, _ = torch.solve(b, A)
            x_exp = solve(A.cpu().numpy(), b.cpu().numpy())
            self.assertEqual(x, x_exp)

        # test against numpy.linalg.solve
        run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6))  # no broadcasting
        run_test((2, 1, 3, 4, 4), (4, 6))  # broadcasting b
        run_test((4, 4), (2, 1, 3, 4, 2))  # broadcasting A
        run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5))
//  [python]
def test_old_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve

        def run_test(A_dims, b_dims):
            A_matrix_size = A_dims[-1]
            A_batch_dims = A_dims[:-2]
            b, A = self.solve_test_helper(A_batch_dims + (A_matrix_size, A_matrix_size), b_dims, device, dtype)
            x, _ = torch.linalg.solve(A, b)
            x_exp = solve(A.cpu().numpy(), b.cpu().numpy())
            self.assertEqual(x, x_exp)

        # test against numpy.linalg.solve
        run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6))  # no broadcasting
        run_test((2, 1, 3, 4, 4), (4, 6))  # broadcasting b
        run_test((4, 4), (2, 1, 3, 4, 2))  # broadcasting A
        run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5))
//  [python]
def convert_theta(mv_preal_locs_ori, MV_Thetas_G):
  ones   = torch.ones(mv_preal_locs_ori.size(0), mv_preal_locs_ori.size(1), mv_preal_locs_ori.size(2), 1, device=mv_preal_locs_ori.device)
  points = torch.cat((mv_preal_locs_ori, ones), dim=-1)
  #trans_points, _ = torch.gesv(points.transpose(-1,-2), MV_Thetas_G)
  trans_points, _ = torch.solve(points.transpose(-1,-2), MV_Thetas_G)
  ok_points = trans_points.transpose(-1,-2)
  return ok_points[..., :2]
//  [python]
def convert_theta(mv_preal_locs_ori, MV_Thetas_G):
  ones   = torch.ones(mv_preal_locs_ori.size(0), mv_preal_locs_ori.size(1), mv_preal_locs_ori.size(2), 1, device=mv_preal_locs_ori.device)
  points = torch.cat((mv_preal_locs_ori, ones), dim=-1)
  #trans_points, _ = torch.gesv(points.transpose(-1,-2), MV_Thetas_G)
  trans_points, _ = torch.linalg.solve(MV_Thetas_G, points.transpose(-1,-2))
  ok_points = trans_points.transpose(-1,-2)
  return ok_points[..., :2]
//  [python]
def M(self, q, eps=1e-4):
        """Returns a function that multiplies the mass matrix M(q) by a vector qdot
        Args:
            q: bs x D Tensor representing the position
        """
        assert q.ndim == 2
        lower_triangular = self.tril_Minv(q)
        assert lower_triangular.ndim == 3

        def M_func(qdot):
            assert qdot.ndim == 2
            qdot = qdot.unsqueeze(-1)
            diag_noise = eps * torch.eye(lower_triangular.size(-1), dtype=qdot.dtype, device=qdot.device)
            # print(f"mass matrix: {(lower_triangular @ lower_triangular.transpose(-2, -1) + diag_noise).mean()}")
            Minv = lower_triangular @ lower_triangular.transpose(-2, -1) + diag_noise
            # print(Minv.mean(0))
            # print(torch.symeig(Minv)[0].mean(0))

            M_times_qdot = torch.solve(qdot, Minv).solution.squeeze(-1)
            return M_times_qdot

        return M_func
//  [python]
def M(self, q, eps=1e-4):
        """Returns a function that multiplies the mass matrix M(q) by a vector qdot
        Args:
            q: bs x D Tensor representing the position
        """
        assert q.ndim == 2
        lower_triangular = self.tril_Minv(q)
        assert lower_triangular.ndim == 3

        def M_func(qdot):
            assert qdot.ndim == 2
            qdot = qdot.unsqueeze(-1)
            diag_noise = eps * torch.eye(lower_triangular.size(-1), dtype=qdot.dtype, device=qdot.device)
            Minv = lower_triangular @ lower_triangular.transpose(-2, -1) + diag_noise
            M_times_qdot = torch.linalg.solve(Minv, qdot).squeeze(-1)
            return M_times_qdot

        return M_func
//  [python]
def forward(self, embeddings, offset_embeddings):
        layer1_act = F.relu(self.layer1(embeddings)) # (num_conj, dim)
        attention = F.softmax(self.layer2(layer1_act), dim=0) # (num_conj, dim)
        #embedding = torch.sum(attention * embeddings, dim=0)
        embeddings = attention*embeddings
        embeddings = torch.unsqueeze(torch.squeeze(embeddings), dim=-2)
        offset_embeddings = torch.unsqueeze(torch.squeeze(offset_embeddings), dim=-2)
        inv_cov1 = torch.matmul(torch.transpose(offset_embeddings[0],-2,-1),offset_embeddings[0])
        inv_cov2 = torch.matmul(torch.transpose(offset_embeddings[-1],-2,-1),offset_embeddings[-1])
        inv_cov_total = inv_cov1+inv_cov2
        mul_term = torch.matmul(inv_cov1,torch.transpose(embeddings[0],-2,-1))+torch.matmul(inv_cov2,torch.transpose(embeddings[-1],-2,-1))
        updated_mean = torch.squeeze(torch.solve(mul_term,inv_cov_total)[0])
        return updated_mean
//  [python]
def forward(self, embeddings, offset_embeddings):
        layer1_act = F.relu(self.layer1(embeddings)) # (num_conj, dim)
        attention = F.softmax(self.layer2(layer1_act), dim=0) # (num_conj, dim)
        #embedding = torch.sum(attention * embeddings, dim=0)
        embeddings = attention*embeddings
        embeddings = torch.unsqueeze(torch.squeeze(embeddings), dim=-2)
        offset_embeddings = torch.unsqueeze(torch.squeeze(offset_embeddings), dim=-2)
        inv_cov1 = torch.matmul(torch.transpose(offset_embeddings[0],-2,-1),offset_embeddings[0])
        inv_cov2 = torch.matmul(torch.transpose(offset_embeddings[-1],-2,-1),offset_embeddings[-1])
        inv_cov_total = inv_cov1+inv_cov2
        mul_term = torch.matmul(inv_cov1,torch.transpose(embeddings[0],-2,-1))+torch.matmul(inv_cov2,torch.transpose(embeddings[-1],-2,-1))
        updated_mean = torch.squeeze(torch.linalg.solve(inv_cov_total,mul_term))
        return updated_mean
//  [python]
def gaussian_wasserstein_commutative(policy: AbstractGaussianPolicy, p: Tuple[ch.Tensor, ch.Tensor],
                                     q: Tuple[ch.Tensor, ch.Tensor], scale_prec=False) -> Tuple[ch.Tensor, ch.Tensor]:
    """
    Compute mean part and cov part of W_2(p || q) with p,q ~ N(y, SS).
    This version DOES assume commutativity of both distributions, i.e. covariance matrices.
    This is less general and assumes both distributions are somewhat close together.
    When scale_prec is true scale both distributions with old precision matrix.
    Args:
        policy: current policy
        p: mean and sqrt of gaussian p
        q: mean and sqrt of gaussian q
        scale_prec: scale objective by old precision matrix.
                    This penalizes directions based on old uncertainty/covariance.

    Returns: mean part of W2, cov part of W2

    """
    mean, sqrt = p
    mean_other, sqrt_other = q

    mean_part = mean_distance(policy, mean, mean_other, sqrt_other, scale_prec)

    cov = policy.covariance(sqrt)
    if scale_prec:
        # cov constraint scaled with precision of old dist
        batch_dim, dim = mean.shape

        identity = ch.eye(dim, dtype=sqrt.dtype, device=sqrt.device)
        sqrt_inv_other = ch.solve(identity, sqrt_other)[0]
        c = sqrt_inv_other @ cov @ sqrt_inv_other

        cov_part = torch_batched_trace(identity + c - 2 * sqrt_inv_other @ sqrt)

    else:
        # W2 objective for cov assuming normal W2 objective for mean
        cov_other = policy.covariance(sqrt_other)
        cov_part = torch_batched_trace(cov_other + cov - 2 * sqrt_other @ sqrt)

    return mean_part, cov_part
//  [python]
def gaussian_wasserstein_commutative(policy: AbstractGaussianPolicy, p: Tuple[ch.Tensor, ch.Tensor],
                                     q: Tuple[ch.Tensor, ch.Tensor], scale_prec=False) -> Tuple[ch.Tensor, ch.Tensor]:
    """
    Compute mean part and cov part of W_2(p || q) with p,q ~ N(y, SS).
    This version DOES assume commutativity of both distributions, i.e. covariance matrices.
    This is less general and assumes both distributions are somewhat close together.
    When scale_prec is true scale both distributions with old precision matrix.
    Args:
        policy: current policy
        p: mean and sqrt of gaussian p
        q: mean and sqrt of gaussian q
        scale_prec: scale objective by old precision matrix.
                    This penalizes directions based on old uncertainty/covariance.

    Returns: mean part of W2, cov part of W2

    """
    mean, sqrt = p
    mean_other, sqrt_other = q

    mean_part = mean_distance(policy, mean, mean_other, sqrt_other, scale_prec)

    cov = policy.covariance(sqrt)
    if scale_prec:
        # cov constraint scaled with precision of old dist
        batch_dim, dim = mean.shape

        identity = ch.eye(dim, dtype=sqrt.dtype, device=sqrt.device)
        sqrt_inv_other = ch.linalg.solve(identity, sqrt_other)
        c = sqrt_inv_other @ cov @ sqrt_inv_other

        cov_part = torch_batched_trace(identity + c - 2 * sqrt_inv_other @ sqrt)

    else:
        # W2 objective for cov assuming normal W2 objective for mean
        cov_other = policy.covariance(sqrt_other)
        cov_part = torch_batched_trace(cov_other + cov - 2 * sqrt_other @ sqrt)

    return mean_part, cov_part
//  [python]
def M(self, q, eps=1e-4):
        """Returns a function that multiplies the mass matrix M(q) by a vector qdot
        Args:
            q: bs x D Tensor representing the position
        """
        assert q.ndim == 2
        lower_triangular = self.tril_Minv(q)
        assert lower_triangular.ndim == 3

        def M_func(qdot):
            assert qdot.ndim == 2
            qdot = qdot.unsqueeze(-1)
            diag_noise = eps * torch.eye(lower_triangular.size(-1), dtype=qdot.dtype, device=qdot.device)
            M_times_qdot = torch.solve(
                    qdot,
                    lower_triangular @ lower_triangular.transpose(-2, -1) + diag_noise
            ).solution.squeeze(-1)
            return M_times_qdot

        return M_func
//  [python]
def M(self, q, eps=1e-4):
        """Returns a function that multiplies the mass matrix M(q) by a vector qdot
        Args:
            q: bs x D Tensor representing the position
        """
        assert q.ndim == 2
        lower_triangular = self.tril_Minv(q)
        assert lower_triangular.ndim == 3

        def M_func(qdot):
            assert qdot.ndim == 2
            qdot = qdot.unsqueeze(-1)
            diag_noise = eps * torch.eye(lower_triangular.size(-1), dtype=qdot.dtype, device=qdot.device)
            M_times_qdot = torch.linalg.solve(
                    lower_triangular @ lower_triangular.transpose(-2, -1) + diag_noise,
                    qdot
            ).squeeze(-1)
            return M_times_qdot

        return M_func
//  [python]
def forward(self, X, Y, w, h):
        """ 计算grid"""
        device = X.device
        grid = torch.ones(1, h, w, 2, device=device)
        grid[:, :, :, 0] = torch.linspace(-1, 1, w)
        grid[:, :, :, 1] = torch.linspace(-1, 1, h)[..., None]
        grid = grid.view(-1, h * w, 2)

        """ 计算W, A"""
        n, k = X.shape[:2]
        device = X.device

        Z = torch.zeros(1, k + 3, 2, device=device)
        P = torch.ones(n, k, 3, device=device)
        L = torch.zeros(n, k + 3, k + 3, device=device)

        eps = 1e-9
        D2 = torch.pow(X[:, :, None, :] - X[:, None, :, :], 2).sum(-1)
        K = D2 * torch.log(D2 + eps)

        P[:, :, 1:] = X
        Z[:, :k, :] = Y
        L[:, :k, :k] = K
        L[:, :k, k:] = P
        L[:, k:, :k] = P.permute(0, 2, 1)

        Q = torch.solve(Z, L)[0]
        W, A = Q[:, :k], Q[:, k:]

        """ 计算U """
        eps = 1e-9
        D2 = torch.pow(grid[:, :, None, :] - X[:, None, :, :], 2).sum(-1)
        U = D2 * torch.log(D2 + eps)

        """ 计算P """
        n, k = grid.shape[:2]
        device = grid.device
        P = torch.ones(n, k, 3, device=device)
        P[:, :, 1:] = grid

        # grid = P @ A + U @ W
        grid = torch.matmul(P, A) + torch.matmul(U, W)
        return grid.view(-1, h, w, 2)
//  [python]
def forward(self, X, Y, w, h):
        """ 计算grid"""
        device = X.device
        grid = torch.ones(1, h, w, 2, device=device)
        grid[:, :, :, 0] = torch.linspace(-1, 1, w)
        grid[:, :, :, 1] = torch.linspace(-1, 1, h)[..., None]
        grid = grid.view(-1, h * w, 2)

        """ 计算W, A"""
        n, k = X.shape[:2]
        device = X.device

        Z = torch.zeros(1, k + 3, 2, device=device)
        P = torch.ones(n, k, 3, device=device)
        L = torch.zeros(n, k + 3, k + 3, device=device)

        eps = 1e-9
        D2 = torch.pow(X[:, :, None, :] - X[:, None, :, :], 2).sum(-1)
        K = D2 * torch.log(D2 + eps)

        P[:, :, 1:] = X
        Z[:, :k, :] = Y
        L[:, :k, :k] = K
        L[:, :k, k:] = P
        L[:, k:, :k] = P.permute(0, 2, 1)

        Q = torch.linalg.solve(L, Z)
        W, A = Q[:, :k], Q[:, k:]

        """ 计算U """
        eps = 1e-9
        D2 = torch.pow(grid[:, :, None, :] - X[:, None, :, :], 2).sum(-1)
        U = D2 * torch.log(D2 + eps)

        """ 计算P """
        n, k = grid.shape[:2]
        device = grid.device
        P = torch.ones(n, k, 3, device=device)
        P[:, :, 1:] = grid

        # grid = P @ A + U @ W
        grid = torch.matmul(P, A) + torch.matmul(U, W)
        return grid.view(-1, h, w, 2)
//  [python]
def solve_affine(self,x, y, w):
        """
        :param x: BxNxD
        :param y: BxNxD
        :param w: BxNx1
        :return:
        """
        # Optimal affine transform: ================================================
        # A = (X^T  @ diag(w) @ X)^-1   @   (X^T  @ diag(w) @ y)
        #    (B,D+1,N)  (B,N,N)  (B,N,D+1)       (B,D+1,N)  (B,N,N)  (B,N,D)
        #
        #   =        Xt_wX    \      Xt_yw
        #           (B,D+1,D+1)       (B,D+1, D)
        # (x, y, z, 1) array to work easily with affine transforms:
        X = torch.cat((x, torch.ones_like(x[:, :, :1])), dim=2)  # (B,N, D+1)

        Xt_wX = X.transpose(2, 1) @ (w * X)  # (B,D+1, N) @ (B,N, D+1) = (B,D+1, D+1)
        Xt_wy = X.transpose(2, 1) @ (w * y)  # (B,D+1, N) @ (B,N, D)   = (B,D+1, D)

        # Affine transformation:
        A = torch.solve(Xt_wy, Xt_wX).solution  # (B,D+1, D)
        return A, X @ A
//  [python]
def solve_affine(self,x, y, w):
        """
        :param x: BxNxD
        :param y: BxNxD
        :param w: BxNx1
        :return:
        """
        # Optimal affine transform: ================================================
        # A = (X^T  @ diag(w) @ X)^-1   @   (X^T  @ diag(w) @ y)
        #    (B,D+1,N)  (B,N,N)  (B,N,D+1)       (B,D+1,N)  (B,N,N)  (B,N,D)
        #
        #   =        Xt_wX    \      Xt_yw
        #           (B,D+1,D+1)       (B,D+1, D)
        # (x, y, z, 1) array to work easily with affine transforms:
        X = torch.cat((x, torch.ones_like(x[:, :, :1])), dim=2)  # (B,N, D+1)

        Xt_wX = X.transpose(2, 1) @ (w * X)  # (B,D+1, N) @ (B,N, D+1) = (B,D+1, D+1)
        Xt_wy = X.transpose(2, 1) @ (w * y)  # (B,D+1, N) @ (B,N, D)   = (B,D+1, D)

        # Affine transformation:
        A = torch.linalg.solve(Xt_wX, Xt_wy)  # (B,D+1, D)
        return A, X @ A
//  [python]
def b_inv(b_mat):
    '''
    code from
    https://stackoverflow.com/questions/46595157/how-to-apply-the-torch-inverse-function-of-pytorch-to-every-sample-in-the-batc
    :param b_mat:
    :return:
    '''
    eye = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    # b_inv, _ = torch.gesv(eye, b_mat) # Updating the PyTorch Version
    try:
        b_inv, _ = torch.solve(eye, b_mat)
    except RuntimeError: # Singularity
        b_inv = torch.pinverse(b_mat)
    return b_inv
//  [python]
def b_inv(b_mat):
    '''
    code from
    https://stackoverflow.com/questions/46595157/how-to-apply-the-torch-inverse-function-of-pytorch-to-every-sample-in-the-batc
    :param b_mat:
    :return:
    '''
    eye = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    # b_inv, _ = torch.gesv(eye, b_mat) # Updating the PyTorch Version
    try:
        b_inv = torch.linalg.solve(b_mat, eye)
    except RuntimeError: # Singularity
        b_inv = torch.pinverse(b_mat)
    return b_inv
//  [python]
def __call__(self, input, *args, **kwargs):
        return_gaussian = kwargs.pop("return_gaussian", True)

        if not self.training and isinstance(input, MultivariateNormal):
            if self.should_return_gaussian_on_eval and return_gaussian:
                targets = args[0]
                if isinstance(targets, list):
                    targets = targets[0].squeeze(-1)
                # noise_cov = self.expected_hessian(input.mean, targets).inverse()
                expected_hessian = self.expected_hessian(input.mean, targets)
                if not isinstance(expected_hessian, torch.Tensor):
                    noise_cov = expected_hessian.inverse()
                else:
                    eye_like_hessian = torch.eye(
                        expected_hessian.shape[-1],
                        device=expected_hessian.device,
                        dtype=expected_hessian.dtype,
                    )
                    noise_cov = torch.solve(
                        eye_like_hessian,
                        lazify(expected_hessian).add_jitter(1e-4).evaluate(),
                    )[0]
                    evals, evecs = noise_cov.symeig(eigenvectors=True)
                    evals[evals <= 0.0] = 0.0
                    noise_cov = evecs.matmul(torch.diag_embed(evals)).matmul(
                        evecs.transpose(-1, -2)
                    )

                return MultivariateNormal(
                    input.mean, input.lazy_covariance_matrix + noise_cov
                )
        return super().__call__(input, *args, **kwargs)
//  [python]
def __call__(self, input, *args, **kwargs):
        return_gaussian = kwargs.pop("return_gaussian", True)

        if not self.training and isinstance(input, MultivariateNormal):
            if self.should_return_gaussian_on_eval and return_gaussian:
                targets = args[0]
                if isinstance(targets, list):
                    targets = targets[0].squeeze(-1)
                # noise_cov = self.expected_hessian(input.mean, targets).inverse()
                expected_hessian = self.expected_hessian(input.mean, targets)
                if not isinstance(expected_hessian, torch.Tensor):
                    noise_cov = expected_hessian.inverse()
                else:
                    eye_like_hessian = torch.eye(
                        expected_hessian.shape[-1],
                        device=expected_hessian.device,
                        dtype=expected_hessian.dtype,
                    )
                    noise_cov = torch.linalg.solve(
                        lazify(expected_hessian).add_jitter(1e-4).evaluate(),
                        eye_like_hessian,
                    )
                    evals, evecs = noise_cov.symeig(eigenvectors=True)
                    evals[evals <= 0.0] = 0.0
                    noise_cov = evecs.matmul(torch.diag_embed(evals)).matmul(
                        evecs.transpose(-1, -2)
                    )

                return MultivariateNormal(
                    input.mean, input.lazy_covariance_matrix + noise_cov
                )
        return super().__call__(input, *args, **kwargs)
//  [python]
def b_inv(b_mat):
    '''
    code from
    https://stackoverflow.com/questions/46595157/how-to-apply-the-torch-inverse-function-of-pytorch-to-every-sample-in-the-batc
    :param b_mat:
    :return:
    '''
    eye = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    try:
        b_inv, _ = torch.solve(eye, b_mat)
    except:
        b_inv = eye
    return b_inv
//  [python]
def b_inv(b_mat):
    '''
    code from
    https://stackoverflow.com/questions/46595157/how-to-apply-the-torch-inverse-function-of-pytorch-to-every-sample-in-the-batc
    :param b_mat:
    :return:
    '''
    eye = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    try:
        b_inv = torch.linalg.solve(b_mat, eye)
    except:
        b_inv = eye
    return b_inv
//  [python]
def _test_inv_matmul(self, rhs, lhs=None, cholesky=False):
        linear_operator = self.create_linear_operator().requires_grad_(True)
        linear_operator_copy = linear_operator.clone().detach_().requires_grad_(True)
        evaluated = self.evaluate_linear_operator(linear_operator_copy)
        evaluated.register_hook(_ensure_symmetric_grad)

        # Create a test right hand side and left hand side
        rhs.requires_grad_(True)
        rhs_copy = rhs.clone().detach().requires_grad_(True)
        if lhs is not None:
            lhs.requires_grad_(True)
            lhs_copy = lhs.clone().detach().requires_grad_(True)

        _wrapped_cg = MagicMock(wraps=linear_cg)
        with patch("linear_operator.utils.linear_cg", new=_wrapped_cg) as linear_cg_mock:
            with settings.max_cholesky_size(math.inf if cholesky else 0), settings.cg_tolerance(1e-4):
                # Perform the inv_matmul
                if lhs is not None:
                    res = linear_operator.inv_matmul(rhs, lhs)
                    actual = lhs_copy @ evaluated.inverse() @ rhs_copy
                else:
                    res = torch.solve(rhs, linear_operator)
                    actual = evaluated.inverse().matmul(rhs_copy)
                self.assertAllClose(res, actual, rtol=0.02, atol=1e-5)

                # Perform backward pass
                grad = torch.randn_like(res)
                res.backward(gradient=grad)
                actual.backward(gradient=grad)
                for arg, arg_copy in zip(linear_operator.representation(), linear_operator_copy.representation()):
                    if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                        self.assertAllClose(arg.grad, arg_copy.grad, rtol=0.03, atol=1e-5)
                self.assertAllClose(rhs.grad, rhs_copy.grad, rtol=0.03, atol=1e-5)
                if lhs is not None:
                    self.assertAllClose(lhs.grad, lhs_copy.grad, rtol=0.03, atol=1e-5)

            # Determine if we've called CG or not
            if not cholesky and self.__class__.should_call_cg:
                self.assertTrue(linear_cg_mock.called)
            else:
                self.assertFalse(linear_cg_mock.called)
//  [python]
def _test_inv_matmul(self, rhs, lhs=None, cholesky=False):
        linear_operator = self.create_linear_operator().requires_grad_(True)
        linear_operator_copy = linear_operator.clone().detach_().requires_grad_(True)
        evaluated = self.evaluate_linear_operator(linear_operator_copy)
        evaluated.register_hook(_ensure_symmetric_grad)

        # Create a test right hand side and left hand side
        rhs.requires_grad_(True)
        rhs_copy = rhs.clone().detach().requires_grad_(True)
        if lhs is not None:
            lhs.requires_grad_(True)
            lhs_copy = lhs.clone().detach().requires_grad_(True)

        _wrapped_cg = MagicMock(wraps=linear_cg)
        with patch("linear_operator.utils.linear_cg", new=_wrapped_cg) as linear_cg_mock:
            with settings.max_cholesky_size(math.inf if cholesky else 0), settings.cg_tolerance(1e-4):
                # Perform the inv_matmul
                if lhs is not None:
                    res = linear_operator.inv_matmul(rhs, lhs)
                    actual = lhs_copy @ evaluated.inverse() @ rhs_copy
                else:
                    res = torch.linalg.solve(linear_operator, rhs)
                    actual = evaluated.inverse().matmul(rhs_copy)
                self.assertAllClose(res, actual, rtol=0.02, atol=1e-5)

                # Perform backward pass
                grad = torch.randn_like(res)
                res.backward(gradient=grad)
                actual.backward(gradient=grad)
                for arg, arg_copy in zip(linear_operator.representation(), linear_operator_copy.representation()):
                    if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
                        self.assertAllClose(arg.grad, arg_copy.grad, rtol=0.03, atol=1e-5)
                self.assertAllClose(rhs.grad, rhs_copy.grad, rtol=0.03, atol=1e-5)
                if lhs is not None:
                    self.assertAllClose(lhs.grad, lhs_copy.grad, rtol=0.03, atol=1e-5)

            # Determine if we've called CG or not
            if not cholesky and self.__class__.should_call_cg:
                self.assertTrue(linear_cg_mock.called)
            else:
                self.assertFalse(linear_cg_mock.called)
//  [python]
def _spin_mean_var(self, x, cav_mean, cav_var):
        """
        Compute spin means and variances from cavity means and variances.

        Note:
            These expressions are obtained from integrating the single-site
            partition function with a multivariate Gaussian prior. You should
            change this function is you want to play around with different
            single-site priors for the spins.
        """
        inv_var = self.spin_prior_inv_var - cav_var
        prefactor = torch.solve(batched_eye_like(inv_var), inv_var).solution
        spin_mean = torch.einsum(
            'i d e, b i d -> b i e', prefactor, (cav_mean + x)
        )
        spin_var = prefactor
        return spin_mean, spin_var
//  [python]
def _spin_mean_var(self, x, cav_mean, cav_var):
        """
        Compute spin means and variances from cavity means and variances.

        Note:
            These expressions are obtained from integrating the single-site
            partition function with a multivariate Gaussian prior. You should
            change this function is you want to play around with different
            single-site priors for the spins.
        """
        inv_var = self.spin_prior_inv_var - cav_var
        prefactor = torch.linalg.solve(batched_eye_like(inv_var), inv_var)
        spin_mean = torch.einsum(
            'i d e, b i d -> b i e', prefactor, (cav_mean + x)
        )
        spin_var = prefactor
        return spin_mean, spin_var
//  [python]
def horo_project_using_one_ideal(submanifold_ideals, x, custom_ideal_direction=None, keep_ambient=True):
    """The first horospherical projection we discussed, currently section 5.3 of the overleaf doc
    Args:
        submanifold_ideals: torch.tensor of shape (sub_dim, dim)
        x: torch.tensor of shape (batch_size, dim)
        custom_ideal_direction (optional): torch.tensor of shape (dim, )
        keep_ambient: boolean

    Returns:
        if keep_ambient == True:
            torch.tensor of shape (batch_size, dim)
        else:
            torch.tensor of shape (batch_size, sub_dim): the rotated projections

    Note:
        custom_ideal_direction, if provided, must be a unit vector in the row span of submanifold_ideals
        submanifold_ideals must have independent rows
        The submanifold we are projecting onto always passes through the origin.
    """

    if custom_ideal_direction is None:
        p = submanifold_ideals[0, :]
        p = p / torch.sqrt(p.dot(p))
    else:
        p = custom_ideal_direction

    eucl_proj_coefs, _ = torch.solve(submanifold_ideals @ x.transpose(0, 1),
                                     submanifold_ideals @ submanifold_ideals.transpose(0, 1))  # (sub_dim, batch_size)
    eucl_projs = eucl_proj_coefs.transpose(0, 1) @ submanifold_ideals  # (batch_size, dim)

    t = torch.sum((p - x) * (p - x), dim=-1) / (
            2 * torch.sum((p - eucl_projs) * (p - x), dim=-1))  # shape (batch_size, )
    t = t.unsqueeze(-1)

    output = 2 * t * eucl_projs + (1 - 2 * t) * p.unsqueeze(0)  # shape (batch_size, dim)
    if keep_ambient:
        return output
    else:
        q, r = torch.qr(submanifold_ideals.transpose(0, 1))  # q.shape (dim, sub_dim)
        return output @ q
//  [python]
def horo_project_using_one_ideal(submanifold_ideals, x, custom_ideal_direction=None, keep_ambient=True):
    """The first horospherical projection we discussed, currently section 5.3 of the overleaf doc
    Args:
        submanifold_ideals: torch.tensor of shape (sub_dim, dim)
        x: torch.tensor of shape (batch_size, dim)
        custom_ideal_direction (optional): torch.tensor of shape (dim, )
        keep_ambient: boolean

    Returns:
        if keep_ambient == True:
            torch.tensor of shape (batch_size, dim)
        else:
            torch.tensor of shape (batch_size, sub_dim): the rotated projections

    Note:
        custom_ideal_direction, if provided, must be a unit vector in the row span of submanifold_ideals
        submanifold_ideals must have independent rows
        The submanifold we are projecting onto always passes through the origin.
    """

    if custom_ideal_direction is None:
        p = submanifold_ideals[0, :]
        p = p / torch.sqrt(p.dot(p))
    else:
        p = custom_ideal_direction

    eucl_proj_coefs, _ = torch.linalg.solve(submanifold_ideals @ submanifold_ideals.transpose(0, 1),
                                     submanifold_ideals @ x.transpose(0, 1))  # (sub_dim, batch_size)
    eucl_projs = eucl_proj_coefs.transpose(0, 1) @ submanifold_ideals  # (batch_size, dim)

    t = torch.sum((p - x) * (p - x), dim=-1) / (
            2 * torch.sum((p - eucl_projs) * (p - x), dim=-1))  # shape (batch_size, )
    t = t.unsqueeze(-1)

    output = 2 * t * eucl_projs + (1 - 2 * t) * p.unsqueeze(0)  # shape (batch_size, dim)
    if keep_ambient:
        return output
    else:
        q, r = torch.qr(submanifold_ideals.transpose(0, 1))  # q.shape (dim, sub_dim)
        return output @ q
//  [python]
def solveAlphas(Pw,control_w):
    N = Pw.shape[0]
    if Pw.is_cuda:
        ones = torch.ones((1,N)).cuda()
    else:
        ones = torch.ones((1,N))
    ph_w = torch.cat([Pw.T,ones],dim=0)
    alphas, LU = torch.solve(ph_w,control_w)

    return alphas
//  [python]
def solveAlphas(Pw,control_w):
    N = Pw.shape[0]
    if Pw.is_cuda:
        ones = torch.ones((1,N)).cuda()
    else:
        ones = torch.ones((1,N))
    ph_w = torch.cat([Pw.T,ones],dim=0)
    alphas = torch.linalg.solve(control_w, ph_w)

    return alphas
//  [python]
def control(self, x, t=None):
        u_dim = self.R.shape[-1]
        with torch.no_grad():
            x_g = self.x_goal
            P = self.Q
            R = self.R * self.dt
            λ = 0.5
            fx = (self.dt * self.model.f_func(x))
            Gx = (self.dt * self.model.g_func(x.unsqueeze(0)).squeeze(0))
            # xp = x + fx + Gx @ u
            # (1-λ) (x + fx + Gx @ u - x_g)ᵀ P (x_g - (x + fx + Gx @ u)) + λ uᵀ R u
            # Quadratic term: uᵀ (λ R + (1-λ)GₓᵀPGₓ) u
            # Linear term   : - (2(1-λ)GₓP(x_g - x - fx)  )ᵀ u
            # Constant term : + (1-λ)(x_g-fx)ᵀP(x_g- x - fx)
            # Minima at u* = ((λ R + (1-λ)GₓᵀPGₓ))⁻¹ ((1-λ)GₓP(x_g - x - fx)  )


            # Quadratic term λ R + (1-λ)Gₓᵀ P Gₓ
            Q = λ * R + (1-λ) * Gx.T @ P @ Gx
            # Linear term - ((1-λ)Gₓ P(x_g - x - fx)  )ᵀ u
            c = (1-λ) * Gx.T @ P @ (x_g - x - fx)
            ugreedy = torch.solve(c.unsqueeze(-1), Q).solution.reshape(-1)
            assert ugreedy.shape[-1] == u_dim
            return ugreedy
//  [python]
def control(self, x, t=None):
        u_dim = self.R.shape[-1]
        with torch.no_grad():
            x_g = self.x_goal
            P = self.Q
            R = self.R * self.dt
            λ = 0.5
            fx = (self.dt * self.model.f_func(x))
            Gx = (self.dt * self.model.g_func(x.unsqueeze(0)).squeeze(0))
            # xp = x + fx + Gx @ u
            # (1-λ) (x + fx + Gx @ u - x_g)ᵀ P (x_g - (x + fx + Gx @ u)) + λ uᵀ R u
            # Quadratic term: uᵀ (λ R + (1-λ)GₓᵀPGₓ) u
            # Linear term   : - (2(1-λ)GₓP(x_g - x - fx)  )ᵀ u
            # Constant term : + (1-λ)(x_g-fx)ᵀP(x_g- x - fx)
            # Minima at u* = ((λ R + (1-λ)GₓᵀPGₓ))⁻¹ ((1-λ)GₓP(x_g - x - fx)  )


            # Quadratic term λ R + (1-λ)Gₓᵀ P Gₓ
            Q = λ * R + (1-λ) * Gx.T @ P @ Gx
            # Linear term - ((1-λ)Gₓ P(x_g - x - fx)  )ᵀ u
            c = (1-λ) * Gx.T @ P @ (x_g - x - fx)
            ugreedy = torch.linalg.solve(Q, c.unsqueeze(-1)).reshape(-1)
            assert ugreedy.shape[-1] == u_dim
            return ugreedy
//  [python]
def batchSolveAlphas(Pw,control_w):
    b = Pw.shape[0]
    if Pw.is_cuda:
        ones = torch.ones((b,1,68)).cuda()
    else:
        ones = torch.ones((b,1,68))

    ph_w = torch.cat([Pw.permute(0,2,1),ones],dim=1)
    alphas, LU = torch.solve(ph_w,control_w)

    return alphas
//  [python]
def batchSolveAlphas(Pw,control_w):
    b = Pw.shape[0]
    if Pw.is_cuda:
        ones = torch.ones((b,1,68)).cuda()
    else:
        ones = torch.ones((b,1,68))

    ph_w = torch.cat([Pw.permute(0,2,1),ones],dim=1)
    alphas = torch.linalg.solve(ph_w,control_w)

    return alphas
//  [python]
def _transp_follow_one(self, x, v, *, u):
        a = self._amat(x, u)
        rhs = v + 1 / 2 * a @ v
        lhs = -1 / 2 * a
        lhs[..., torch.arange(a.shape[-2]), torch.arange(x.shape[-2])] += 1
        qv, _ = torch.solve(rhs, lhs)
        return qv
//  [python]
def _transp_follow_one(self, x, v, *, u):
        a = self._amat(x, u)
        rhs = v + 1 / 2 * a @ v
        lhs = -1 / 2 * a
        lhs[..., torch.arange(a.shape[-2]), torch.arange(x.shape[-2])] += 1
        qv = torch.linalg.solve(lhs, rhs)
        return qv
//  [python]
def test_least_squares(self):
        set_seed(243)
        m, n = 100, 20

        A = cp.Parameter((m, n))
        b = cp.Parameter(m)
        x = cp.Variable(n)
        obj = cp.sum_squares(A@x - b) + cp.sum_squares(x)
        prob = cp.Problem(cp.Minimize(obj))
        prob_th = CvxpyLayer(prob, [A, b], [x])

        A_th = torch.randn(m, n).double().requires_grad_()
        b_th = torch.randn(m).double().requires_grad_()

        x = prob_th(A_th, b_th, solver_args={"eps": 1e-10})[0]

        def lstsq(
            A,
            b): return torch.solve(
            (A_th.t() @ b_th).unsqueeze(1),
            A_th.t() @ A_th +
            torch.eye(n).double())[0]
        x_lstsq = lstsq(A_th, b_th)

        grad_A_cvxpy, grad_b_cvxpy = grad(x.sum(), [A_th, b_th])
        grad_A_lstsq, grad_b_lstsq = grad(x_lstsq.sum(), [A_th, b_th])

        self.assertAlmostEqual(
            torch.norm(
                grad_A_cvxpy -
                grad_A_lstsq).item(),
            0.0)
        self.assertAlmostEqual(
            torch.norm(
                grad_b_cvxpy -
                grad_b_lstsq).item(),
            0.0)
//  [python]
def test_least_squares(self):
        set_seed(243)
        m, n = 100, 20

        A = cp.Parameter((m, n))
        b = cp.Parameter(m)
        x = cp.Variable(n)
        obj = cp.sum_squares(A@x - b) + cp.sum_squares(x)
        prob = cp.Problem(cp.Minimize(obj))
        prob_th = CvxpyLayer(prob, [A, b], [x])

        A_th = torch.randn(m, n).double().requires_grad_()
        b_th = torch.randn(m).double().requires_grad_()

        x = prob_th(A_th, b_th, solver_args={"eps": 1e-10})[0]

        def lstsq(
            A,
            b): return torch.linalg.solve(
            A_th.t() @ A_th +
            torch.eye(n).double(),
            (A_th.t() @ b_th).unsqueeze(1))
        x_lstsq = lstsq(A_th, b_th)

        grad_A_cvxpy, grad_b_cvxpy = grad(x.sum(), [A_th, b_th])
        grad_A_lstsq, grad_b_lstsq = grad(x_lstsq.sum(), [A_th, b_th])

        self.assertAlmostEqual(
            torch.norm(
                grad_A_cvxpy -
                grad_A_lstsq).item(),
            0.0)
        self.assertAlmostEqual(
            torch.norm(
                grad_b_cvxpy -
                grad_b_lstsq).item(),
            0.0)
//  [python]
def apply_affine2point(points, theta, shape):
  assert points.size(0) == 3, 'invalid points shape : {:}'.format(points.size())
  with torch.no_grad():
    ok_points = points[2,:] == 1
    assert torch.sum(ok_points).item() > 0, 'there is no visiable point'
    points[:2,:] = normalize_points(shape, points[:2,:])

    norm_trans_points = ok_points.unsqueeze(0).repeat(3, 1).float()

    #trans_points, ___ = torch.gesv(points[:, ok_points], theta) # removed in PyTorch 1.2.0
    trans_points, ___ = torch.solve(points[:, ok_points], theta)
    

    norm_trans_points[:, ok_points] = trans_points
    
  return norm_trans_points