jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def log_partition_function(self):
        # By Kirchoff's matrix-tree theorem, the partition function is the
        # determinant of a truncated version of the graph Laplacian matrix. We
        # use a Cholesky decomposition to compute the log determinant.
        # See https://en.wikipedia.org/wiki/Kirchhoff%27s_theorem
        V = self.num_vertices
        v1, v2 = make_complete_graph(V).unbind(0)
        logits = self.edge_logits.new_full((V, V), -math.inf)
        logits[v1, v2] = self.edge_logits
        logits[v2, v1] = self.edge_logits
        log_diag = logits.logsumexp(-1)
        # Numerically stabilize so that laplacian has 1's on the diagonal.
        shift = 0.5 * log_diag
        laplacian = torch.eye(V) - (logits - shift - shift[:, None]).exp()
        truncated = laplacian[:-1, :-1]
        try:
            import gpytorch

            log_det = gpytorch.lazy.NonLazyTensor(truncated).logdet()
        except ImportError:
            log_det = torch.cholesky(truncated).diag().log().sum() * 2
        return log_det + log_diag[:-1].sum()
//  [python]
def log_partition_function(self):
        # By Kirchoff's matrix-tree theorem, the partition function is the
        # determinant of a truncated version of the graph Laplacian matrix. We
        # use a Cholesky decomposition to compute the log determinant.
        # See https://en.wikipedia.org/wiki/Kirchhoff%27s_theorem
        V = self.num_vertices
        v1, v2 = make_complete_graph(V).unbind(0)
        logits = self.edge_logits.new_full((V, V), -math.inf)
        logits[v1, v2] = self.edge_logits
        logits[v2, v1] = self.edge_logits
        log_diag = logits.logsumexp(-1)
        # Numerically stabilize so that laplacian has 1's on the diagonal.
        shift = 0.5 * log_diag
        laplacian = torch.eye(V) - (logits - shift - shift[:, None]).exp()
        truncated = laplacian[:-1, :-1]
        try:
            import gpytorch

            log_det = gpytorch.lazy.NonLazyTensor(truncated).logdet()
        except ImportError:
            log_det = torch.linalg.cholesky(truncated).diag().log().sum() * 2
        return log_det + log_diag[:-1].sum()
//  [python]
def log_likelihood(x, mu, sigma, pi):

    chol = torch.cholesky(sigma)  # [nk, nc, nc]
    log_det = chol.diagonal(0, -1, -2).log().sum(-1).mul(2)  # [nk]
    chol = chol.inverse()

    z = (x - mu[:, :, None])  # [nk, nc, nv]
    z = z.movedim(-1, 0).unsqueeze(-1)  # [nv, nk, nc, 1]
    z = torch.matmul(chol, z)  # [nv, nk, nc, 1]
    z = z.squeeze(-1).square_().sum(-1)  # [nv, nk]
    z += log_det
    z *= -0.5
    z += pi.log()
    z = z.T  # [nk, nv]

    return z
//  [python]
def log_likelihood(x, mu, sigma, pi):

    chol = torch.linalg.cholesky(sigma)  # [nk, nc, nc]
    log_det = chol.diagonal(0, -1, -2).log().sum(-1).mul(2)  # [nk]
    chol = chol.inverse()

    z = (x - mu[:, :, None])  # [nk, nc, nv]
    z = z.movedim(-1, 0).unsqueeze(-1)  # [nv, nk, nc, 1]
    z = torch.matmul(chol, z)  # [nv, nk, nc, 1]
    z = z.squeeze(-1).square_().sum(-1)  # [nv, nk]
    z += log_det
    z *= -0.5
    z += pi.log()
    z = z.T  # [nk, nv]

    return z
//  [python]
def model(self):
        self.set_mode("model")

        N = self.X.size(0)
        Kff = self.kernel(self.X).contiguous()
        Kff.view(-1)[:: N + 1] += self.jitter  # add jitter to the diagonal
        Lff = torch.cholesky(Kff)

        zero_loc = self.X.new_zeros(self.f_loc.shape)
        if self.whiten:
            identity = eye_like(self.X, N)
            pyro.sample(
                self._pyro_get_fullname("f"),
                dist.MultivariateNormal(zero_loc, scale_tril=identity).to_event(
                    zero_loc.dim() - 1
                ),
            )
            f_scale_tril = Lff.matmul(self.f_scale_tril)
            f_loc = Lff.matmul(self.f_loc.unsqueeze(-1)).squeeze(-1)
        else:
            pyro.sample(
                self._pyro_get_fullname("f"),
                dist.MultivariateNormal(zero_loc, scale_tril=Lff).to_event(
                    zero_loc.dim() - 1
                ),
            )
            f_scale_tril = self.f_scale_tril
            f_loc = self.f_loc

        f_loc = f_loc + self.mean_function(self.X)
        f_var = f_scale_tril.pow(2).sum(dim=-1)
        if self.y is None:
            return f_loc, f_var
        else:
            return self.likelihood(f_loc, f_var, self.y)
//  [python]
def model(self):
        self.set_mode("model")

        N = self.X.size(0)
        Kff = self.kernel(self.X).contiguous()
        Kff.view(-1)[:: N + 1] += self.jitter  # add jitter to the diagonal
        Lff = torch.linalg.cholesky(Kff)

        zero_loc = self.X.new_zeros(self.f_loc.shape)
        if self.whiten:
            identity = eye_like(self.X, N)
            pyro.sample(
                self._pyro_get_fullname("f"),
                dist.MultivariateNormal(zero_loc, scale_tril=identity).to_event(
                    zero_loc.dim() - 1
                ),
            )
            f_scale_tril = Lff.matmul(self.f_scale_tril)
            f_loc = Lff.matmul(self.f_loc.unsqueeze(-1)).squeeze(-1)
        else:
            pyro.sample(
                self._pyro_get_fullname("f"),
                dist.MultivariateNormal(zero_loc, scale_tril=Lff).to_event(
                    zero_loc.dim() - 1
                ),
            )
            f_scale_tril = self.f_scale_tril
            f_loc = self.f_loc

        f_loc = f_loc + self.mean_function(self.X)
        f_var = f_scale_tril.pow(2).sum(dim=-1)
        if self.y is None:
            return f_loc, f_var
        else:
            return self.likelihood(f_loc, f_var, self.y)
//  [python]
def fit(self, dataset: torch.Tensor) -> None:
        """Fit a KDE model to the input dataset.

        Args:
          dataset (torch.Tensor): Input dataset.

        Returns:
            None
        """
        num_samples, dimension = dataset.shape

        # compute scott's bandwidth factor
        factor = num_samples ** (-1 / (dimension + 4))

        cov_mat = self.cov(dataset.T)
        inv_cov_mat = torch.linalg.inv(cov_mat)
        inv_cov = inv_cov_mat / factor**2

        # transform data to account for bandwidth
        bw_transform = torch.cholesky(inv_cov)
        dataset = torch.matmul(dataset, bw_transform)

        #
        norm = torch.prod(torch.diag(bw_transform))
        norm *= math.pow((2 * math.pi), (-dimension / 2))

        self.bw_transform = bw_transform
        self.dataset = dataset
        self.norm = norm
//  [python]
def fit(self, dataset: torch.Tensor) -> None:
        """Fit a KDE model to the input dataset.

        Args:
          dataset (torch.Tensor): Input dataset.

        Returns:
            None
        """
        num_samples, dimension = dataset.shape

        # compute scott's bandwidth factor
        factor = num_samples ** (-1 / (dimension + 4))

        cov_mat = self.cov(dataset.T)
        inv_cov_mat = torch.linalg.inv(cov_mat)
        inv_cov = inv_cov_mat / factor**2

        # transform data to account for bandwidth
        bw_transform = torch.linalg.cholesky(inv_cov)
        dataset = torch.matmul(dataset, bw_transform)

        #
        norm = torch.prod(torch.diag(bw_transform))
        norm *= math.pow((2 * math.pi), (-dimension / 2))

        self.bw_transform = bw_transform
        self.dataset = dataset
        self.norm = norm
//  [python]
def _estimate_posterior(self, x_test):
        """Estimate Posterior (mu, cov) of GP conditioned on data

        x_test : test points
        """
        cov11 = self._kernel(self.x_data, self.x_data)
        cov12 = self._kernel(self.x_data, x_test)
        cov22 = self._kernel(x_test, x_test)
        term1, _ = torch.solve(cov12, cov11)
        mu_2_1 = torch.mm(term1.T, self.y_data.view(-1, 1)).view(-1)
        cov_2_1 = cov22 - torch.mm(term1.T, cov12)
        return mu_2_1, cov_2_1
//  [python]
def _estimate_posterior(self, x_test):
        """Estimate Posterior (mu, cov) of GP conditioned on data

        x_test : test points
        """
        cov11 = self._kernel(self.x_data, self.x_data)
        cov12 = self._kernel(self.x_data, x_test)
        cov22 = self._kernel(x_test, x_test)
        term1 = torch.linalg.solve(cov11, cov12)
        mu_2_1 = torch.mm(term1.T, self.y_data.view(-1, 1)).view(-1)
        cov_2_1 = cov22 - torch.mm(term1.T, cov12)
        return mu_2_1, cov_2_1
//  [python]
def _solve_P_Q(U, V):
    P = U + V
    Q = -U + V
    return torch.solve(P, Q)[0]
//  [python]
def _solve_P_Q(U, V):
    P = U + V
    Q = -U + V
    return torch.linalg.solve(P, Q)
//  [python]
def b_inv(b_mat):
    ''' Performs batch matrix inversion.

    Arguments:
        b_mat: the batch of matrices that should be inverted
    '''

    eye = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    # b_inv, _ = torch.gesv(eye, b_mat)
    b_inv, _ = torch.solve(eye, b_mat)
    return b_inv
//  [python]
def b_inv(b_mat):
    ''' Performs batch matrix inversion.

    Arguments:
        b_mat: the batch of matrices that should be inverted
    '''

    eye = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    b_inv, _ = torch.linalg.solve(eye, b_mat)
    return b_inv
//  [python]
def init_bias(self, ts, us, zupts):
        """Init IMU bias and orientation with first measurements"""
        if zupts[:self.N_init].sum() == self.N_init:
            N_init = self.N_init
        else:
            N_init = torch.where(zupts[:self.N_init] == 0)[0][0].item()
            cprint('Bias initialized with only {} samples'.format(N_init),
                'yellow')
        u = us[:N_init].mean(dim=0)
        gravity = -u[3:6]
        self.Rot = self.SO3.from_2vectors(gravity, self.g)
        self.b_omega = u[:3]
        self.b_acc = self.Rot.t().mv(self.g) - gravity

        self.P = self.P0.clone()
        # init covariance
        H = u.new_zeros(6, 15)
        H[:3, 9:12] = self.Id3
        H[3:6, :3] = self.Rot.t().mm(self.Wg)
        H[3:6, 12:15] = -self.Id3
        R = self.Q[:6, :6]
        for i in range(N_init):
            S = axat(H, self.P) + R
            Kt, _ = torch.solve(self.P.mm(H.t()).t(), S)
            K = Kt.t()
            I_KH = self.IdP - K.mm(H)
            self.P = axat(I_KH, self.P.clone()) + axat(K, R)
        self.P = (self.P + self.P.t()).clone()/2
//  [python]
def init_bias(self, ts, us, zupts):
        """Init IMU bias and orientation with first measurements"""
        if zupts[:self.N_init].sum() == self.N_init:
            N_init = self.N_init
        else:
            N_init = torch.where(zupts[:self.N_init] == 0)[0][0].item()
            cprint('Bias initialized with only {} samples'.format(N_init),
                'yellow')
        u = us[:N_init].mean(dim=0)
        gravity = -u[3:6]
        self.Rot = self.SO3.from_2vectors(gravity, self.g)
        self.b_omega = u[:3]
        self.b_acc = self.Rot.t().mv(self.g) - gravity

        self.P = self.P0.clone()
        # init covariance
        H = u.new_zeros(6, 15)
        H[:3, 9:12] = self.Id3
        H[3:6, :3] = self.Rot.t().mm(self.Wg)
        H[3:6, 12:15] = -self.Id3
        R = self.Q[:6, :6]
        for i in range(N_init):
            S = axat(H, self.P) + R
            Kt, _ = torch.linalg.solve(S, self.P.mm(H.t()).t())
            K = Kt.t()
            I_KH = self.IdP - K.mm(H)
            self.P = axat(I_KH, self.P.clone()) + axat(K, R)
        self.P = (self.P + self.P.t()).clone()/2
//  [python]
def forward(self, x):
        B, D, H, W = x.shape
        N = H * W
        identity = torch.eye(N).cuda()
        # reshape x, s.t. we can use the gmp formulation as a global pooling operation
        x = x.view(B, D, N)
        x = x.permute(0, 2, 1)
        # compute the linear kernel
        K = torch.bmm(x, x.permute(0, 2, 1))
        # solve the linear system (K + lambda * I) * alpha = ones
        A = K + self.lamb * identity
        o = torch.ones(B, N, 1).cuda()
        #alphas, _ = torch.gesv(o, A) # tested using pytorch 1.0.1
        alphas, _ = torch.solve(o, A) # tested using pytorch 1.2.0
        alphas = alphas.view(B, 1, -1)        
        xi = torch.bmm(alphas, x)
        xi = xi.view(B, -1)
        return xi
//  [python]
def forward(self, x):
        B, D, H, W = x.shape
        N = H * W
        identity = torch.eye(N).cuda()
        # reshape x, s.t. we can use the gmp formulation as a global pooling operation
        x = x.view(B, D, N)
        x = x.permute(0, 2, 1)
        # compute the linear kernel
        K = torch.bmm(x, x.permute(0, 2, 1))
        # solve the linear system (K + lambda * I) * alpha = ones
        A = K + self.lamb * identity
        o = torch.ones(B, N, 1).cuda()
        alphas, _ = torch.linalg.solve(A, o) 
        alphas = alphas.view(B, 1, -1)        
        xi = torch.bmm(alphas, x)
        xi = xi.view(B, -1)
        return xi
//  [python]
def update(self, i, u, cov, zupt):
        H = self.H
        if zupt == 1:
            z = 1
            cov = cov[0] * cov.new_ones(3)
        else:
            z = 0

        H[:3, 3:6] = self.Rot.t()
        H[0, 3:6] *= z

        self.r = torch.cat((- self.Rot.t().mv(self.v),
            u[:3] - self.b_omega,
            u[3:6] - self.b_acc + self.Rot.t().mv(self.g)))
        self.r[0] *= z

        z *= self.r[3:6].norm() < self.max_omega_norm
        z *= self.r[3:6].abs().max() < self.max_omega
        z *= self.r[6:9].norm() < self.max_acc_norm
        z *= self.r[6:9].abs().max() < self.max_acc
        self.r[3:9] *= z

        H[3:6, 9:12] = z*self.Id3
        H[6:9, 12:15] = z*self.Id3
        H[6:9, :3] = -z*self.Rot.t().mm(self.Wg)

        R = torch.diag(torch.cat((cov,
            self.zupt_omega_cov,
            self.zupt_acc_cov), 0))
        S = axat(H, self.P) + R
        Kt, _ = torch.solve(self.P.mm(H.t()).t(), S)
        K = Kt.t()
        self.xi = K.mv(self.r)
        self.state_update(i)
        self.covariance_update(i, K, H, R)
//  [python]
def update(self, i, u, cov, zupt):
        H = self.H
        if zupt == 1:
            z = 1
            cov = cov[0] * cov.new_ones(3)
        else:
            z = 0

        H[:3, 3:6] = self.Rot.t()
        H[0, 3:6] *= z

        self.r = torch.cat((- self.Rot.t().mv(self.v),
            u[:3] - self.b_omega,
            u[3:6] - self.b_acc + self.Rot.t().mv(self.g)))
        self.r[0] *= z

        z *= self.r[3:6].norm() < self.max_omega_norm
        z *= self.r[3:6].abs().max() < self.max_omega
        z *= self.r[6:9].norm() < self.max_acc_norm
        z *= self.r[6:9].abs().max() < self.max_acc
        self.r[3:9] *= z

        H[3:6, 9:12] = z*self.Id3
        H[6:9, 12:15] = z*self.Id3
        H[6:9, :3] = -z*self.Rot.t().mm(self.Wg)

        R = torch.diag(torch.cat((cov,
            self.zupt_omega_cov,
            self.zupt_acc_cov), 0))
        S = axat(H, self.P) + R
        Kt, _ = torch.linalg.solve(S, self.P.mm(H.t()).t())
        K = Kt.t()
        self.xi = K.mv(self.r)
        self.state_update(i)
        self.covariance_update(i, K, H, R)
//  [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, k), device, dtype)
            x = torch.solve(b, A)[0]
            self.assertEqual(b, A.mm(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, k), device, dtype)
            x = torch.linalg.solve(A, b)
            self.assertEqual(b, A.mm(x))
//  [python]
def test_old_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value
        A = random_fullrank_matrix_distinct_singular_value(2, 2, dtype=dtype).to(device).permute(1, 0, 2)
        b = torch.randn(2, 2, 2, dtype=dtype, device=device).permute(2, 1, 0)
        x, _ = torch.solve(b, A)
        x_exp = solve(A.cpu().numpy(), b.cpu().numpy())
        self.assertEqual(x, x_exp)
//  [python]
def test_old_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value
        A = random_fullrank_matrix_distinct_singular_value(2, 2, dtype=dtype).to(device).permute(1, 0, 2)
        b = torch.randn(2, 2, 2, dtype=dtype, device=device).permute(2, 1, 0)
        x, _ = torch.linalg.solve(A, b)
        x_exp = solve(A.cpu().numpy(), b.cpu().numpy())
        self.assertEqual(x, x_exp)
//  [python]
def test_old_solve_batched(self, device, dtype):
        def solve_batch_helper(A_dims, b_dims):
            b, A = self.solve_test_helper(A_dims, b_dims, device, dtype)
            x_exp_list = []
            for i in range(b_dims[0]):
                x_exp_list.append(torch.solve(b[i], A[i])[0])
            x_exp = torch.stack(x_exp_list)  # Stacked output
            x_act = torch.solve(b, A)[0]  # Actual output
            self.assertEqual(x_exp, x_act)  # Equality check
            Ax = torch.matmul(A, x_act)
            self.assertEqual(b, Ax)

        for batchsize in [1, 3, 4]:
            solve_batch_helper((5, batchsize), (batchsize, 5, 10))
//  [python]
def test_old_solve_batched(self, device, dtype):
        def solve_batch_helper(A_dims, b_dims):
            b, A = self.solve_test_helper(A_dims, b_dims, device, dtype)
            x_exp_list = []
            for i in range(b_dims[0]):
                x_exp_list.append(torch.linalg.solve(A[i], b[i]))
            x_exp = torch.stack(x_exp_list)  # Stacked output
            x_act = torch.linalg.solve(A, b)  # Actual output
            self.assertEqual(x_exp, x_act)  # Equality check
            Ax = torch.matmul(A, x_act)
            self.assertEqual(b, Ax)

        for batchsize in [1, 3, 4]:
            solve_batch_helper((5, batchsize), (batchsize, 5, 10))
//  [python]
def test_old_solve_batched_many_batches(self, device, dtype):
        for A_dims, b_dims in zip([(5, 256, 256), (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([(5, 256, 256), (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 unproject2(self, p2d, Kdist, Mdist, distCoeffs):
        pn2d, _ = torch.linalg.solve(Kdist, torch.cat((p2d[:,:2].T, torch.ones((1, p2d.shape[0])).cuda().double()), axis=0))
        pn2d = pn2d.T

        if vis_check_flag:
            print(f'pn2d.mean() {pn2d.mean()}')

        k = torch.cat((distCoeffs[:5], torch.zeros(12-5).cuda().double()), axis=0)
        x0 = pn2d[:,0]
        y0 = pn2d[:,1]
        x = x0
        y = y0

        for iter in range(5):
            r2 = x*x + y*y
            icdist = (1 + ((k[7]*r2 + k[6])*r2 + k[5])*r2)/(1 + ((k[4]*r2 + k[1])*r2 + k[0])*r2)
            deltaX = 2*k[2]*x*y + k[3]*(r2 + 2*x*x)+ k[8]*r2+k[9]*r2*r2
            deltaY = k[2]*(r2 + 2*y*y) + 2*k[3]*x*y+ k[10]*r2+k[11]*r2*r2
            x = (x0 - deltaX)*icdist
            y = (y0 - deltaY)*icdist
        
        pn2d = torch.cat((x.unsqueeze(-1), y.unsqueeze(-1)), axis=1)
        
        if vis_check_flag:
            print(f'pn2d.mean() {pn2d.mean()}')
            print(f'p2d.mean() {p2d.mean()}')
            print(f'(p2d[:,2:3]*0.001).mean() {(p2d[:,2:3]*0.001).mean()}')
            print(f'(pn2d * (p2d[:,2:3]*0.001)).mean() {(pn2d * (p2d[:,2:3]*0.001)).mean()}')
            print(f'pn2d {pn2d}')
        

        p3d = torch.cat([pn2d * (p2d[:,2:3]*0.001), p2d[:,2:3]*0.001], axis=1)
        
        p3d = torch.matmul(torch.inverse(torch.cat([Mdist, torch.DoubleTensor([[0,0,0,1]]).cuda()], dim=0)), torch.cat((p3d, torch.ones((p3d.shape[0], 1)).cuda().double()), axis=1).T).T

        if vis_check_flag:
            print(f'p3d.mean() {p3d.mean()}')

        p3d = p3d[:,:3]
        return p3d
//  [python]
def forward(self, X):
        # (I + X)(I - X)^{-1}
        return torch.solve(self.Id + X, self.Id - X).solution
//  [python]
def forward(self, X):
        # (I + X)(I - X)^{-1}
        return torch.linalg.solve(self.Id - X, self.Id + X)
//  [python]
def cayley_map(X):
    n = X.size(0)
    Id = torch.eye(n, dtype=X.dtype, device=X.device)
    return torch.solve(Id - X, Id + X)[0]
//  [python]
def cayley_map(X):
    n = X.size(0)
    Id = torch.eye(n, dtype=X.dtype, device=X.device)
    return torch.linalg.solve(Id - X, Id + X)
//  [python]
def forward(self, X):
        # Assume X skew-symmetric
        # (I + X)(I - X)^{-1}
        return torch.solve(self.Id + X, self.Id - X).solution
//  [python]
def forward(self, X):
        # Assume X skew-symmetric
        # (I + X)(I - X)^{-1}
        return torch.linalg.solve(self.Id - X, self.Id + X)
//  [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_matrix_size,) + A_batch_dims, 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_matrix_size,) + A_batch_dims, 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 binv(b_mat):
    """
    Computes an inverse of each matrix in the batch.
    Pytorch 0.4.1 does not support batched matrix inverse.
    Hence, we are solving AX=I.

    Parameters:
      b_mat:  a (n_batch, n, n) Tensor.
    Returns: a (n_batch, n, n) Tensor.
    """

    id_matrix = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat).to(b_mat.device)
    b_inv, _ = torch.solve(id_matrix, b_mat)

    return b_inv
//  [python]
def binv(b_mat):
    """
    Computes an inverse of each matrix in the batch.
    Pytorch 0.4.1 does not support batched matrix inverse.
    Hence, we are solving AX=I.

    Parameters:
      b_mat:  a (n_batch, n, n) Tensor.
    Returns: a (n_batch, n, n) Tensor.
    """

    id_matrix = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat).to(b_mat.device)
    b_inv = torch.linalg.solve(b_mat, id_matrix)

    return b_inv
//  [python]
def test_solve_methods_arg_device(self, device):
        for b_device, A_device in itertools.product(['cpu', device], repeat=2):
            if b_device == A_device:
                continue

            b = torch.randn(3, 1, device=b_device)
            A = torch.randn(3, 3, device=A_device)
            err_str = "Expected b and A to be on the same device"
            with self.assertRaisesRegex(RuntimeError, err_str):
                torch.solve(b, A)

            with self.assertRaisesRegex(RuntimeError, err_str):
                torch.cholesky_solve(b, A)

            with self.assertRaisesRegex(RuntimeError, err_str):
                torch.triangular_solve(b, A)

            # b and A have to be modified to match accepted inputs sizes for lu_solve
            b = b.unsqueeze(0)
            A = A.unsqueeze(0)
            with self.assertRaisesRegex(RuntimeError, err_str):
                torch.lu_solve(b, A, torch.rand(A.shape[:-1], device=A_device).int())

            # This checks if a suitable error message is thrown
            # when LU output and pivots are on the same device
            with self.assertRaisesRegex(RuntimeError,
                                        "Expected LU_pivots and LU_data to be on the same device"):
                torch.lu_solve(b, A, torch.rand(A.shape[:-1], device=b_device).int())
//  [python]
def test_solve_methods_arg_device(self, device):
        for b_device, A_device in itertools.product(['cpu', device], repeat=2):
            if b_device == A_device:
                continue

            b = torch.randn(3, 1, device=b_device)
            A = torch.randn(3, 3, device=A_device)
            err_str = "Expected b and A to be on the same device"
            with self.assertRaisesRegex(RuntimeError, err_str):
                torch.linalg.solve(A, b)

            with self.assertRaisesRegex(RuntimeError, err_str):
                torch.cholesky_solve(b, A)

            with self.assertRaisesRegex(RuntimeError, err_str):
                torch.triangular_solve(b, A)

            # b and A have to be modified to match accepted inputs sizes for lu_solve
            b = b.unsqueeze(0)
            A = A.unsqueeze(0)
            with self.assertRaisesRegex(RuntimeError, err_str):
                torch.lu_solve(b, A, torch.rand(A.shape[:-1], device=A_device).int())

            # This checks if a suitable error message is thrown
            # when LU output and pivots are on the same device
            with self.assertRaisesRegex(RuntimeError,
                                        "Expected LU_pivots and LU_data to be on the same device"):
                torch.lu_solve(b, A, torch.rand(A.shape[:-1], device=b_device).int())
//  [python]
def right_inverse(self, A):
        # Assume A orthogonal
        # See https://en.wikipedia.org/wiki/Cayley_transform#Matrix_map
        # (X - I)(X + I)^{-1}
        return torch.solve(X - self.Id, self.Id + X).solution
//  [python]
def right_inverse(self, A):
        # Assume A orthogonal
        # See https://en.wikipedia.org/wiki/Cayley_transform#Matrix_map
        # (X - I)(X + I)^{-1}
        return torch.linalg.solve(self.Id + X, X - self.Id)
//  [python]
def fit(self, features: torch.Tensor, targets: torch.Tensor, prior_precision: float = 1):
        """
        Learn weights from data using MAP inference
        Args:
            features: (Npoints x Nfeatures) tensor with training features
            targets: (Npoints,) tensor
            prior_precision: Precision of an isotropic Gaussian prior (also known as ridge regulariser)
        Returns:
            None
        """
        assert len(targets) == features.shape[0]
        assert len(features.shape) == 2

        self.posterior_prec = (
            features.T @ features
            + torch.eye(features.shape[1], dtype=features.dtype, device=features.device) * prior_precision
        )

        self.w = torch.solve(features.T, self.posterior_prec)[0] @ targets
//  [python]
def fit(self, features: torch.Tensor, targets: torch.Tensor, prior_precision: float = 1):
        """
        Learn weights from data using MAP inference
        Args:
            features: (Npoints x Nfeatures) tensor with training features
            targets: (Npoints,) tensor
            prior_precision: Precision of an isotropic Gaussian prior (also known as ridge regulariser)
        Returns:
            None
        """
        assert len(targets) == features.shape[0]
        assert len(features.shape) == 2

        self.posterior_prec = (
            features.T @ features
            + torch.eye(features.shape[1], dtype=features.dtype, device=features.device) * prior_precision
        )

        self.w = torch.linalg.solve(self.posterior_prec, features.T) @ targets
//  [python]
def predict(self, features: torch.Tensor, compute_covariance=False):
        """
        Make predictions
        Args:
            features: (Npoints x Nfeatures) tensor containing test features
            compute_covariance: whether to compute and return a covariance matrix over test targets
        Returns:
            pred_mu: a (Npoints) tensor containing predicted values for the test points
            pred_cov: A (Npoints x Npoints) covariance matrix if compute_covariance is True, else None
        """
        assert len(self.w) == features.shape[1]
        assert self.w is not None, "model must be fit before it can make predictions"

        pred_mu = features @ self.w
        if compute_covariance:
            pred_cov = features @ torch.solve(features.T, self.posterior_prec)[0]
        else:
            pred_cov = None

        return pred_mu, pred_cov
//  [python]
def predict(self, features: torch.Tensor, compute_covariance=False):
        """
        Make predictions
        Args:
            features: (Npoints x Nfeatures) tensor containing test features
            compute_covariance: whether to compute and return a covariance matrix over test targets
        Returns:
            pred_mu: a (Npoints) tensor containing predicted values for the test points
            pred_cov: A (Npoints x Npoints) covariance matrix if compute_covariance is True, else None
        """
        assert len(self.w) == features.shape[1]
        assert self.w is not None, "model must be fit before it can make predictions"

        pred_mu = features @ self.w
        if compute_covariance:
            pred_cov = features @ torch.linalg.solve(self.posterior_prec, features.T)
        else:
            pred_cov = None

        return pred_mu, pred_cov
//  [python]
def solve(self, v, regul=1e-8, impl='solve'):
        """
        solves v = Ax in x
        """
        # TODO: test
        if impl == 'solve':
            # TODO: reuse LU decomposition once it is computed
            inv_v, _ = torch.solve(v.get_flat_representation().view(-1, 1),
                                   self.data +
                                   regul * torch.eye(self.size(0),
                                                     device=self.data.device))
            return PVector(v.layer_collection, vector_repr=inv_v[:, 0])
        elif impl == 'eigendecomposition':
            v_eigenbasis = self.project_to_diag(v)
            inv_v_eigenbasis = v_eigenbasis / (self.evals + regul)
            return self.project_from_diag(inv_v_eigenbasis)
        else:
            raise NotImplementedError
//  [python]
def solve(self, v, regul=1e-8, impl='solve'):
        """
        solves v = Ax in x
        """
        # TODO: test
        if impl == 'solve':
            # TODO: reuse LU decomposition once it is computed
            inv_v, _ = torch.linalg.solve(self.data + regul * torch.eye(self.size(0), device=self.data.device), v.get_flat_representation().view(-1, 1))
            return PVector(v.layer_collection, vector_repr=inv_v[:, 0])
        elif impl == 'eigendecomposition':
            v_eigenbasis = self.project_to_diag(v)
            inv_v_eigenbasis = v_eigenbasis / (self.evals + regul)
            return self.project_from_diag(inv_v_eigenbasis)
        else:
            raise NotImplementedError
//  [python]
def project_onto_constraints(G,z,tol=1e-5):
    """inputs [Graph (n,E)] [z (bs,2,n,d)] ouput [Pz (bs,2,n,d)]
       Runs several iterations of Newton-Raphson to minimize the constraint violation """
    bs,_,n,d = z.shape
    violation = np.inf
    with torch.no_grad():
        i=0
        while violation>tol:
            Phi = rigid_Phi(G,z[:,0],z[:,1]) # (bs,2,C)
            DPhi = rigid_DPhi(G,z[:,0],z[:,1]) # (bs,2,n,d,2,C)
            violation = (Phi**2).mean().sqrt()
            J = DPhi.reshape(bs,2*n*d,-1).permute(0,2,1)
            if J.shape[-2]<J.shape[-1]:
                #Jinv = torch.pinverse(J)
                #diff = -(Jinv@Phi.reshape(bs,-1,1)).reshape(*z.shape)
                diff = (J.permute(0,2,1)@torch.solve(-Phi.reshape(bs,-1,1),J@J.permute(0,2,1))[0]).reshape(*z.shape)
            else: #cry
                print(J.shape)
            #print(violation)
            scale = (z**2).mean().sqrt()
            z += diff.clamp(min=-scale/2,max=scale/2)
            i+=1
            if i>500: raise OverflowError("Newton-Raphson Constraint projection failed to converge")
        #print(f"converged in {i} iterations")
    return z
//  [python]
def project_onto_constraints(G,z,tol=1e-5):
    """inputs [Graph (n,E)] [z (bs,2,n,d)] ouput [Pz (bs,2,n,d)]
       Runs several iterations of Newton-Raphson to minimize the constraint violation """
    bs,_,n,d = z.shape
    violation = np.inf
    with torch.no_grad():
        i=0
        while violation>tol:
            Phi = rigid_Phi(G,z[:,0],z[:,1]) # (bs,2,C)
            DPhi = rigid_DPhi(G,z[:,0],z[:,1]) # (bs,2,n,d,2,C)
            violation = (Phi**2).mean().sqrt()
            J = DPhi.reshape(bs,2*n*d,-1).permute(0,2,1)
            if J.shape[-2]<J.shape[-1]:
                #Jinv = torch.pinverse(J)
                #diff = -(Jinv@Phi.reshape(bs,-1,1)).reshape(*z.shape)
                diff = (J.permute(0,2,1)@torch.linalg.solve(J@J.permute(0,2,1),-Phi.reshape(bs,-1,1))).reshape(*z.shape)
            else: #cry
                print(J.shape)
            #print(violation)
            scale = (z**2).mean().sqrt()
            z += diff.clamp(min=-scale/2,max=scale/2)
            i+=1
            if i>500: raise OverflowError("Newton-Raphson Constraint projection failed to converge")
        #print(f"converged in {i} iterations")
    return z
//  [python]
def solve(self, vs, regul=1e-8):
        vs_dict = vs.get_dict_representation()
        out_dict = dict()
        for layer_id, layer in self.generator.layer_collection.layers.items():
            v = vs_dict[layer_id][0].view(-1)
            if layer.bias is not None:
                v = torch.cat([v, vs_dict[layer_id][1].view(-1)])
            block = self.data[layer_id]

            inv_v, _ = torch.solve(v.view(-1, 1),
                                   block +
                                   regul * torch.eye(block.size(0),
                                                     device=block.device))
            inv_v_tuple = (inv_v[:layer.weight.numel()]
                           .view(*layer.weight.size),)
            if layer.bias is not None:
                inv_v_tuple = (inv_v_tuple[0],
                               inv_v[layer.weight.numel():]
                               .view(*layer.bias.size),)

            out_dict[layer_id] = inv_v_tuple
        return PVector(layer_collection=vs.layer_collection,
                       dict_repr=out_dict)
//  [python]
def solve(self, vs, regul=1e-8):
        vs_dict = vs.get_dict_representation()
        out_dict = dict()
        for layer_id, layer in self.generator.layer_collection.layers.items():
            v = vs_dict[layer_id][0].view(-1)
            if layer.bias is not None:
                v = torch.cat([v, vs_dict[layer_id][1].view(-1)])
            block = self.data[layer_id]

            inv_v = torch.linalg.solve(block + regul * torch.eye(block.size(0), device=block.device), v.view(-1, 1))
            inv_v_tuple = (inv_v[:layer.weight.numel()]
                           .view(*layer.weight.size),)
            if layer.bias is not None:
                inv_v_tuple = (inv_v_tuple[0],
                               inv_v[layer.weight.numel():]
                               .view(*layer.bias.size),)

            out_dict[layer_id] = inv_v_tuple
        return PVector(layer_collection=vs.layer_collection,
                       dict_repr=out_dict)
//  [python]
def test_old_solve_errors_and_warnings(self, device, dtype):
        # dtypes should be safely castable
        a = torch.eye(2, dtype=dtype, device=device)
        b = torch.randn(2, 1, dtype=dtype, device=device)
        out = torch.empty(0, dtype=torch.int, device=device)
        lu = torch.empty(0, dtype=dtype, device=device)
        with self.assertRaisesRegex(RuntimeError, "but got solution with dtype Int"):
            torch.solve(b, a, out=(out, lu))

        out = torch.empty(0, dtype=dtype, device=device)
        lu = torch.empty(0, dtype=torch.int, device=device)
        with self.assertRaisesRegex(RuntimeError, "but got lu with dtype Int"):
            torch.solve(b, a, out=(out, lu))

        # device should match
        if torch.cuda.is_available():
            wrong_device = 'cpu' if self.device_type != 'cpu' else 'cuda'
            out = torch.empty(0, dtype=dtype, device=wrong_device)
            lu = torch.empty_like(a)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.solve(b, a, out=(out, lu))
            out = torch.empty(0, dtype=dtype, device=device)
            lu = torch.empty_like(a).to(wrong_device)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.solve(b, a, out=(out, lu))
//  [python]
def test_old_solve_errors_and_warnings(self, device, dtype):
        # dtypes should be safely castable
        a = torch.eye(2, dtype=dtype, device=device)
        b = torch.randn(2, 1, dtype=dtype, device=device)
        out = torch.empty(0, dtype=torch.int, device=device)
        lu = torch.empty(0, dtype=dtype, device=device)
        with self.assertRaisesRegex(RuntimeError, "but got solution with dtype Int"):
            torch.linalg.solve(a, b, out=(out, lu))

        out = torch.empty(0, dtype=dtype, device=device)
        lu = torch.empty(0, dtype=torch.int, device=device)
        with self.assertRaisesRegex(RuntimeError, "but got lu with dtype Int"):
            torch.linalg.solve(a, b, out=(out, lu))

        # device should match
        if torch.cuda.is_available():
            wrong_device = 'cpu' if self.device_type != 'cpu' else 'cuda'
            out = torch.empty(0, dtype=dtype, device=wrong_device)
            lu = torch.empty_like(a)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.linalg.solve(a, b, out=(out, lu))
            out = torch.empty(0, dtype=dtype, device=device)
            lu = torch.empty_like(a).to(wrong_device)
            with self.assertRaisesRegex(RuntimeError, "tensors to be on the same device"):
                torch.linalg.solve(a, b, out=(out, lu))
//  [python]
def test_solve_methods_arg_device(self, device):
        for b_device, A_device in itertools.product(['cpu', device], repeat=2):
            if b_device == A_device:
                continue

            b = torch.randn(3, 1, device=b_device)
            A = torch.randn(3, 3, device=A_device)

            # solve and cholesky_solve goes through generic backend dispatch and hit kernel specific device check first
            # triangular_solve goes through specific backend dispatch (CPU/CUDA) and hit auto-generated device check first
            generic_backend_dispatch_err_str = "Expected b and A to be on the same device"
            specific_backend_dispatch_err_str = "Expected all tensors to be on the same device"
            with self.assertRaisesRegex(RuntimeError, generic_backend_dispatch_err_str):
                torch.solve(b, A)

            with self.assertRaisesRegex(RuntimeError, generic_backend_dispatch_err_str):
                torch.cholesky_solve(b, A)

            with self.assertRaisesRegex(RuntimeError, specific_backend_dispatch_err_str):
                torch.triangular_solve(b, A)

            # b and A have to be modified to match accepted inputs sizes for lu_solve
            b = b.unsqueeze(0)
            A = A.unsqueeze(0)
            with self.assertRaisesRegex(RuntimeError, generic_backend_dispatch_err_str):
                torch.lu_solve(b, A, torch.rand(A.shape[:-1], device=A_device).int())

            # This checks if a suitable error message is thrown
            # when LU output and pivots are on the same device
            with self.assertRaisesRegex(RuntimeError,
                                        "Expected LU_pivots and LU_data to be on the same device"):
                torch.lu_solve(b, A, torch.rand(A.shape[:-1], device=b_device).int())
//  [python]
def test_solve_methods_arg_device(self, device):
        for b_device, A_device in itertools.product(['cpu', device], repeat=2):
            if b_device == A_device:
                continue

            b = torch.randn(3, 1, device=b_device)
            A = torch.randn(3, 3, device=A_device)

            # solve and cholesky_solve goes through generic backend dispatch and hit kernel specific device check first
            # triangular_solve goes through specific backend dispatch (CPU/CUDA) and hit auto-generated device check first
            generic_backend_dispatch_err_str = "Expected b and A to be on the same device"
            specific_backend_dispatch_err_str = "Expected all tensors to be on the same device"
            with self.assertRaisesRegex(RuntimeError, generic_backend_dispatch_err_str):
                torch.linalg.solve(A, b)

            with self.assertRaisesRegex(RuntimeError, generic_backend_dispatch_err_str):
                torch.cholesky_solve(b, A)

            with self.assertRaisesRegex(RuntimeError, specific_backend_dispatch_err_str):
                torch.triangular_solve(b, A)

            # b and A have to be modified to match accepted inputs sizes for lu_solve
            b = b.unsqueeze(0)
            A = A.unsqueeze(0)
            with self.assertRaisesRegex(RuntimeError, generic_backend_dispatch_err_str):
                torch.lu_solve(b, A, torch.rand(A.shape[:-1], device=A_device).int())

            # This checks if a suitable error message is thrown
            # when LU output and pivots are on the same device
            with self.assertRaisesRegex(RuntimeError,
                                        "Expected LU_pivots and LU_data to be on the same device"):
                torch.lu_solve(b, A, torch.rand(A.shape[:-1], device=b_device).int())
//  [python]
def solve(self, v, regul=1e-8, impl='solve'):
        # TODO: test
        if impl == 'solve':
            # TODO: reuse LU decomposition once it is computed
            inv_v, _ = torch.solve(v.get_flat_representation().view(-1, 1),
                                   self.data +
                                   regul * torch.eye(self.size(0),
                                                     device=self.data.device))
            return PVector(v.layer_collection, vector_repr=inv_v[:, 0])
        elif impl == 'eigendecomposition':
            v_eigenbasis = self.project_to_diag(v)
            inv_v_eigenbasis = v_eigenbasis / (self.evals + regul)
            return self.project_from_diag(inv_v_eigenbasis)
        else:
            raise NotImplementedError
//  [python]
def solve(self, v, regul=1e-8, impl='solve'):
        # TODO: test
        if impl == 'solve':
            # TODO: reuse LU decomposition once it is computed
            inv_v, _ = torch.linalg.solve(self.data + regul * torch.eye(self.size(0), device=self.data.device), v.get_flat_representation().view(-1, 1))
            return PVector(v.layer_collection, vector_repr=inv_v[:, 0])
        elif impl == 'eigendecomposition':
            v_eigenbasis = self.project_to_diag(v)
            inv_v_eigenbasis = v_eigenbasis / (self.evals + regul)
            return self.project_from_diag(inv_v_eigenbasis)
        else:
            raise NotImplementedError
//  [python]
def retr(self, X, G):
        """
        Retract to positive definite
        """
        X_inv_G, _ = torch.solve(G, X)
        return multisym(X + G + .5 * multiprod(G, X_inv_G))
//  [python]
def retr(self, X, G):
        """
        Retract to positive definite
        """
        X_inv_G = torch.linalg.solve(X, G)
        return multisym(X + G + .5 * multiprod(G, X_inv_G))
//  [python]
def inner(self, X, G1, G2):
        G1s, _ = torch.solve(G1, X)
        G2s, _ = torch.solve(G2, X)
        return torch.sum(G1s * G2s)
//  [python]
def inner(self, X, G1, G2):
        G1s = torch.linalg.solve(X, G1)
        G2s = torch.linalg.solve(X, G2)
        return torch.sum(G1s * G2s)
//  [python]
def torch_expm(A):
    n_A = A.shape[0]
    A_fro = torch.sqrt(A.abs().pow(2).sum(dim=(1, 2), keepdim=True))

    # Scaling step
    maxnorm = torch.tensor([5.371920351148152], dtype=A.dtype, device=A.device)
    zero = torch.tensor([0.0], dtype=A.dtype, device=A.device)
    n_squarings = torch.max(zero, torch.ceil(torch_log2(A_fro / maxnorm)))
    A_scaled = A / 2.0 ** n_squarings
    n_squarings = n_squarings.flatten().type(torch.int64)

    # Pade 13 approximation
    U, V = torch_pade13(A_scaled)
    P = U + V
    Q = -U + V
    R, _ = torch.solve(P, Q)

    # Unsquaring step
    res = [R]
    for i in range(int(n_squarings.max())):
        res.append(res[-1].matmul(res[-1]))
    R = torch.stack(res)
    expmA = R[n_squarings, torch.arange(n_A)]
    return expmA[0]
//  [python]
def torch_expm(A):
    n_A = A.shape[0]
    A_fro = torch.sqrt(A.abs().pow(2).sum(dim=(1, 2), keepdim=True))

    # Scaling step
    maxnorm = torch.tensor([5.371920351148152], dtype=A.dtype, device=A.device)
    zero = torch.tensor([0.0], dtype=A.dtype, device=A.device)
    n_squarings = torch.max(zero, torch.ceil(torch_log2(A_fro / maxnorm)))
    A_scaled = A / 2.0 ** n_squarings
    n_squarings = n_squarings.flatten().type(torch.int64)

    # Pade 13 approximation
    U, V = torch_pade13(A_scaled)
    P = U + V
    Q = -U + V
    R, _ = torch.linalg.solve(P, Q)

    # Unsquaring step
    res = [R]
    for i in range(int(n_squarings.max())):
        res.append(res[-1].matmul(res[-1]))
    R = torch.stack(res)
    expmA = R[n_squarings, torch.arange(n_A)]
    return expmA[0]
//  [python]
def frame2euler(frame_pos_vel):
    """ Inverts euler2frame, converts from frame+derivatives-> euler+derivatives
        input: (bs,2,3,3)
        output: (bs,2,3)"""
    R,Rdot = frame_pos_vel.permute(1,0,3,2) #(bs,3,3)
    omega = uncross_matrix(R.permute(0,2,1)@Rdot) #angular velocity in body frame Omega = RTRdot
    angles = torch.from_numpy(np.ascontiguousarray(Rotation.from_matrix(R.data.cpu().numpy()).as_euler('ZXZ'))).to(R.device,R.dtype)
    eulerdot = torch.solve(omega.unsqueeze(-1),eulerdot2omega(angles))[0].squeeze(-1)
    return torch.stack([angles,eulerdot],dim=1)
//  [python]
def frame2euler(frame_pos_vel):
    """ Inverts euler2frame, converts from frame+derivatives-> euler+derivatives
        input: (bs,2,3,3)
        output: (bs,2,3)"""
    R,Rdot = frame_pos_vel.permute(1,0,3,2) #(bs,3,3)
    omega = uncross_matrix(R.permute(0,2,1)@Rdot) #angular velocity in body frame Omega = RTRdot
    angles = torch.from_numpy(np.ascontiguousarray(Rotation.from_matrix(R.data.cpu().numpy()).as_euler('ZXZ'))).to(R.device,R.dtype)
    eulerdot = torch.linalg.solve(eulerdot2omega(angles),omega.unsqueeze(-1))[0].squeeze(-1)
    return torch.stack([angles,eulerdot],dim=1)
//  [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.solve(v, g_reg)
            solve_a, _ = torch.solve(solve_g.t(), a_reg)
            solve_a = solve_a.t()
            if layer.bias is None:
                solve_tuple = (solve_a.view(*sw),)
            else:
                solve_tuple = (solve_a[:, :-1].contiguous().view(*sw),
                               solve_a[:, -1].contiguous())
            out_dict[layer_id] = solve_tuple
        return PVector(layer_collection=vs.layer_collection,
                       dict_repr=out_dict)