jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_old_eig_invalid_input(self, device, dtype):
        # test invalid input
        self.assertRaisesRegex(
            RuntimeError,
            'input should be 2 dimensional',
            lambda: torch.eig(torch.ones((2))))
        self.assertRaisesRegex(
            RuntimeError,
            'input should be square',
            lambda: torch.eig(torch.ones((2, 3))))
        self.assertRaisesRegex(
            RuntimeError,
            'input should not contain infs or NaNs',
            lambda: torch.eig(np.inf * torch.ones((2, 2))))
        self.assertRaisesRegex(
            RuntimeError,
            'input should not contain infs or NaNs',
            lambda: torch.eig(np.nan * torch.ones((2, 2))))
//  [python]
def test_old_eig_invalid_input(self, device, dtype):
        # test invalid input
        self.assertRaisesRegex(
            RuntimeError,
            'input should be 2 dimensional',
            lambda: torch.linalg.eig(torch.ones((2))))
        self.assertRaisesRegex(
            RuntimeError,
            'input should be square',
            lambda: torch.linalg.eig(torch.ones((2, 3))))
        self.assertRaisesRegex(
            RuntimeError,
            'input should not contain infs or NaNs',
            lambda: torch.linalg.eig(np.inf * torch.ones((2, 2))))
        self.assertRaisesRegex(
            RuntimeError,
            'input should not contain infs or NaNs',
            lambda: torch.linalg.eig(np.nan * torch.ones((2, 2))))
//  [python]
def test_eig_out(self, device, dtype):
        # the out version of torch.eig needs to be tested manually: we can't
        # use the "test_out=True" parameter to tensor_op_tests because the
        # signature is irregular (since we have *two* output vectors)
        t = torch.randn(10, 10, dtype=dtype, device=device)
        evals, evecs = torch.eig(t, eigenvectors=True)
        #
        # check that the out= version computes the same values as the normal one
        out_evals = torch.empty_like(evals)
        out_evecs = torch.empty_like(evecs)
        evals2, evecs2 = torch.eig(t, eigenvectors=True, out=(out_evals, out_evecs))
        # check that the out tensors were used in-place
        self.assertEqual(evals2.data_ptr(), out_evals.data_ptr())
        self.assertEqual(evecs2.data_ptr(), out_evecs.data_ptr())
        # check that the result is the same as the non-out version
        self.assertEqual(evals, out_evals)
        self.assertEqual(evecs, out_evecs)
        #
        # check what happens in the eigenvectors=False case
        out_evals = torch.empty_like(evals)
        out_evecs = torch.tensor([1, 2, 3], dtype=dtype, device=device)
        evals2, evecs2 = torch.eig(t, eigenvectors=False, out=(out_evals, out_evecs))
        # check that the out_evals was used in-place
        self.assertEqual(evals2.data_ptr(), out_evals.data_ptr())
        self.assertEqual(evals, out_evals)
        # check that out_evecs was NOT touched at all
        assert out_evecs.tolist() == [1, 2, 3]
        #
        # check that we complain if we pass an out vector of the wrong dtype
        wrong_out = torch.empty((0, 0), dtype=int)
        with self.assertRaisesRegex(RuntimeError, r"Expected .* but got .*"):
            torch.eig(t, eigenvectors=True, out=(wrong_out, out_evecs))
        with self.assertRaisesRegex(RuntimeError, r"Expected .* but got .*"):
            torch.eig(t, eigenvectors=True, out=(out_evals, wrong_out))
//  [python]
def test_eig_out(self, device, dtype):
        # the out version of torch.linalg.eig needs to be tested manually: we can't
        # use the "test_out=True" parameter to tensor_op_tests because the
        # signature is irregular (since we have *two* output vectors)
        t = torch.randn(10, 10, dtype=dtype, device=device)
        evals, evecs = torch.linalg.eig(t)
        #
        # check that the out= version computes the same values as the normal one
        out_evals = torch.empty_like(evals)
        out_evecs = torch.empty_like(evecs)
        evals2, evecs2 = torch.linalg.eig(t, out=(out_evals, out_evecs))
        # check that the out tensors were used in-place
        self.assertEqual(evals2.data_ptr(), out_evals.data_ptr())
        self.assertEqual(evecs2.data_ptr(), out_evecs.data_ptr())
        # check that the result is the same as the non-out version
        self.assertEqual(evals, out_evals)
        self.assertEqual(evecs, out_evecs)
        #
        # check what happens in the eigenvectors=False case
        out_evals = torch.empty_like(evals)
        out_evecs = torch.tensor([1, 2, 3], dtype=dtype, device=device)
        evals2, evecs2 = torch.linalg.eig(t, out=(out_evals, out_evecs))
        # check that the out_evals was used in-place
        self.assertEqual(evals2.data_ptr(), out_evals.data_ptr())
        self.assertEqual(evals, out_evals)
        # check that out_evecs was NOT touched at all
        assert out_evecs.tolist() == [1, 2, 3]
        #
        # check that we complain if we pass an out vector of the wrong dtype
        wrong_out = torch.empty((0, 0), dtype=int)
        with self.assertRaisesRegex(RuntimeError, r"Expected .* but got .*"):
            torch.linalg.eig(t, out=(wrong_out, out_evecs))
        with self.assertRaisesRegex(RuntimeError, r"Expected .* but got .*"):
            torch.linalg.eig(t, out=(out_evals, wrong_out))
//  [python]
def test_old_eig_non_contiguous(self, device, dtype):
        X = torch.randn(4, 4, dtype=dtype, device=device)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, 2, dtype=dtype, device=device)[:, 1]
        v = torch.zeros(4, 2, 4, dtype=dtype, device=device)[:, 1]
        self.assertFalse(v.is_contiguous(), 'V is contiguous')
        self.assertFalse(e.is_contiguous(), 'E is contiguous')
        torch.eig(X, True, out=(e, v))
        Xhat = np.matmul(np.matmul(v.cpu(), torch.diag(e.cpu().select(1, 0))), v.t().cpu())
        if dtype is torch.float:
            atol = 1e-7
            rtol = 1e-5
        else:
            atol = 1e-8
            rtol = 0
        self.assertEqual(X, Xhat, atol=atol, rtol=rtol, msg='VeV\' wrong')
//  [python]
def test_old_eig_non_contiguous(self, device, dtype):
        X = torch.randn(4, 4, dtype=dtype, device=device)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, 2, dtype=dtype, device=device)[:, 1]
        v = torch.zeros(4, 2, 4, dtype=dtype, device=device)[:, 1]
        self.assertFalse(v.is_contiguous(), 'V is contiguous')
        self.assertFalse(e.is_contiguous(), 'E is contiguous')
        e, v = torch.linalg.eig(X)
        Xhat = np.matmul(np.matmul(v.cpu(), torch.diag(e.cpu().select(1, 0))), v.t().cpu())
        if dtype is torch.float:
            atol = 1e-7
            rtol = 1e-5
        else:
            atol = 1e-8
            rtol = 0
        self.assertEqual(X, Xhat, atol=atol, rtol=rtol, msg='VeV\' wrong')
//  [python]
def test_old_eig_reuse(self, device, dtype):
        X = torch.randn(4, 4, dtype=dtype, device=device)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, dtype=dtype, device=device)
        v = torch.zeros(4, 4, dtype=dtype, device=device)
        torch.eig(X, True, out=(e, v))
        Xhat = np.matmul(np.matmul(v.cpu(), torch.diag(e.select(1, 0)).cpu()), v.t().cpu())
        if dtype is torch.float:
            atol = 1e-7
            rtol = 1e-5
        else:
            atol = 1e-8
            rtol = 0
        self.assertEqual(X, Xhat, atol=atol, rtol=rtol, msg='VeV\' wrong')
        self.assertTrue(v.is_contiguous(), 'V is not contiguous')

        torch.eig(X, True, out=(e, v))
        Xhat = np.matmul(v.cpu(), np.matmul(e.select(1, 0).diag().cpu(), v.t().cpu()))
        self.assertEqual(X, Xhat, atol=atol, rtol=rtol, msg='VeV\' wrong')
        self.assertTrue(v.is_contiguous(), 'V is not contiguous')
//  [python]
def test_old_eig_reuse(self, device, dtype):
        X = torch.randn(4, 4, dtype=dtype, device=device)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, dtype=dtype, device=device)
        v = torch.zeros(4, 4, dtype=dtype, device=device)
        torch.linalg.eig(X, out=(e, v))
        Xhat = np.matmul(np.matmul(v.cpu(), torch.diag(e.select(1, 0)).cpu()), v.t().cpu())
        if dtype is torch.float:
            atol = 1e-7
            rtol = 1e-5
        else:
            atol = 1e-8
            rtol = 0
        self.assertEqual(X, Xhat, atol=atol, rtol=rtol, msg='VeV\' wrong')
        self.assertTrue(v.is_contiguous(), 'V is not contiguous')

        torch.linalg.eig(X, out=(e, v))
        Xhat = np.matmul(v.cpu(), np.matmul(e.select(1, 0).diag().cpu(), v.t().cpu()))
        self.assertEqual(X, Xhat, atol=atol, rtol=rtol, msg='VeV\' wrong')
        self.assertTrue(v.is_contiguous(), 'V is not contiguous')
//  [python]
def is_psd(X, tol=1e-12):
        """Check if given matrix is PSD"""
        evals, evecs = torch.eig(X, eigenvectors=False)
        return torch.all(evals >= 0 - tol)
//  [python]
def is_psd(X, tol=1e-12):
        """Check if given matrix is PSD"""
        evals = torch.linalg.eigvals(X)
        return torch.all(evals.real >= 0 - tol)
//  [python]
def test_old_eig_basic(self, device, dtype):
        a = torch.tensor([[1.96, 0.00, 0.00, 0.00, 0.00],
                          [-6.49, 3.80, 0.00, 0.00, 0.00],
                          [-0.47, -6.39, 4.17, 0.00, 0.00],
                          [-7.20, 1.50, -1.51, 5.70, 0.00],
                          [-0.65, -6.34, 2.67, 1.80, -7.10]],
                         dtype=dtype, device=device).t()
        e = torch.eig(a)[0]
        ee, vv = torch.eig(a, True)
        te = torch.tensor((), dtype=dtype, device=device)
        tv = torch.tensor((), dtype=dtype, device=device)
        eee, vvv = torch.eig(a, True, out=(te, tv))
        self.assertEqual(e, ee, atol=1e-12, rtol=0)
        self.assertEqual(ee, eee, atol=1e-12, rtol=0)
        self.assertEqual(ee, te, atol=1e-12, rtol=0)
        self.assertEqual(vv, vvv, atol=1e-12, rtol=0)
        self.assertEqual(vv, tv, atol=1e-12, rtol=0)
        #
        # compare with numpy
        np_e, np_v = np.linalg.eig(a.cpu().numpy())
        if dtype.is_complex:
            self.assertEqual(ee, np_e)
        else:
            # np_e.shape == (n, 2), where each column contain the real and
            # imaginary parts of the result
            self.assertEqual(ee[:, 0], np_e)  # real part
            self.assertEqual(ee[:, 1], torch.zeros(ee.shape[0], dtype=dtype))  # imaginary part
        self.assertEqual(vv, np_v)
//  [python]
def test_old_eig_basic(self, device, dtype):
        a = torch.tensor([[1.96, 0.00, 0.00, 0.00, 0.00],
                          [-6.49, 3.80, 0.00, 0.00, 0.00],
                          [-0.47, -6.39, 4.17, 0.00, 0.00],
                          [-7.20, 1.50, -1.51, 5.70, 0.00],
                          [-0.65, -6.34, 2.67, 1.80, -7.10]],
                         dtype=dtype, device=device).t()
        e = torch.linalg.eig(a)[0]
        ee, vv = torch.linalg.eig(a, True)
        te = torch.tensor((), dtype=dtype, device=device)
        tv = torch.tensor((), dtype=dtype, device=device)
        eee, vvv = torch.linalg.eig(a, True, out=(te, tv))
        self.assertEqual(e, ee, atol=1e-12, rtol=0)
        self.assertEqual(ee, eee, atol=1e-12, rtol=0)
        self.assertEqual(ee, te, atol=1e-12, rtol=0)
        self.assertEqual(vv, vvv, atol=1e-12, rtol=0)
        self.assertEqual(vv, tv, atol=1e-12, rtol=0)
        #
        # compare with numpy
        np_e, np_v = np.linalg.eig(a.cpu().numpy())
        if dtype.is_complex:
            self.assertEqual(ee, np_e)
        else:
            # np_e.shape == (n, 2), where each column contain the real and
            # imaginary parts of the result
            self.assertEqual(ee[:, 0], np_e)  # real part
            self.assertEqual(ee[:, 1], torch.zeros(ee.shape[0], dtype=dtype))  # imaginary part
        self.assertEqual(vv, np_v)
//  [python]
def test_eig(self):
        a = torch.Tensor(((1.96, 0.00, 0.00, 0.00, 0.00),
                          (-6.49, 3.80, 0.00, 0.00, 0.00),
                          (-0.47, -6.39, 4.17, 0.00, 0.00),
                          (-7.20, 1.50, -1.51, 5.70, 0.00),
                          (-0.65, -6.34, 2.67, 1.80, -7.10))).t().contiguous()
        e = torch.eig(a)[0]
        ee, vv = torch.eig(a, True)
        te = torch.Tensor()
        tv = torch.Tensor()
        eee, vvv = torch.eig(a, True, out=(te, tv))
        self.assertEqual(e, ee, 1e-12)
        self.assertEqual(ee, eee, 1e-12)
        self.assertEqual(ee, te, 1e-12)
        self.assertEqual(vv, vvv, 1e-12)
        self.assertEqual(vv, tv, 1e-12)

        # test reuse
        X = torch.randn(4, 4)
        X = torch.mm(X.t(), X)
        e, v = torch.zeros(4, 2), torch.zeros(4, 4)
        torch.eig(X, True, out=(e, v))
        Xhat = torch.mm(torch.mm(v, torch.diag(e.select(1, 0))), v.t())
        self.assertEqual(X, Xhat, 1e-8, 'VeV\' wrong')
        self.assertFalse(v.is_contiguous(), 'V is contiguous')

        torch.eig(X, True, out=(e, v))
        Xhat = torch.mm(v, torch.mm(e.select(1, 0).diag(), v.t()))
        self.assertEqual(X, Xhat, 1e-8, 'VeV\' wrong')
        self.assertFalse(v.is_contiguous(), 'V is contiguous')

        # test non-contiguous
        X = torch.randn(4, 4)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, 2)[:, 1]
        v = torch.zeros(4, 2, 4)[:, 1]
        self.assertFalse(v.is_contiguous(), 'V is contiguous')
        self.assertFalse(e.is_contiguous(), 'E is contiguous')
        torch.eig(X, True, out=(e, v))
        Xhat = torch.mm(torch.mm(v, torch.diag(e.select(1, 0))), v.t())
        self.assertEqual(X, Xhat, 1e-8, 'VeV\' wrong')
//  [python]
def test_eig(self):
        a = torch.Tensor(((1.96, 0.00, 0.00, 0.00, 0.00),
                          (-6.49, 3.80, 0.00, 0.00, 0.00),
                          (-0.47, -6.39, 4.17, 0.00, 0.00),
                          (-7.20, 1.50, -1.51, 5.70, 0.00),
                          (-0.65, -6.34, 2.67, 1.80, -7.10))).t().contiguous()
        e = torch.linalg.eig(a)[0]
        ee, vv = torch.linalg.eig(a, True)
        te = torch.Tensor()
        tv = torch.Tensor()
        eee, vvv = torch.linalg.eig(a, True, out=(te, tv))
        self.assertEqual(e, ee, 1e-12)
        self.assertEqual(ee, eee, 1e-12)
        self.assertEqual(ee, te, 1e-12)
        self.assertEqual(vv, vvv, 1e-12)
        self.assertEqual(vv, tv, 1e-12)

        # test reuse
        X = torch.randn(4, 4)
        X = torch.mm(X.t(), X)
        e, v = torch.zeros(4, 2), torch.zeros(4, 4)
        torch.linalg.eig(X, True, out=(e, v))
        Xhat = torch.mm(torch.mm(v, torch.diag(e.select(1, 0))), v.t())
        self.assertEqual(X, Xhat, 1e-8, 'VeV\' wrong')
        self.assertFalse(v.is_contiguous(), 'V is contiguous')

        torch.linalg.eig(X, True, out=(e, v))
        Xhat = torch.mm(v, torch.mm(e.select(1, 0).diag(), v.t()))
        self.assertEqual(X, Xhat, 1e-8, 'VeV\' wrong')
        self.assertFalse(v.is_contiguous(), 'V is contiguous')

        # test non-contiguous
        X = torch.randn(4, 4)
        X = torch.mm(X.t(), X)
        e = torch.zeros(4, 2, 2)[:, 1]
        v = torch.zeros(4, 2, 4)[:, 1]
        self.assertFalse(v.is_contiguous(), 'V is contiguous')
        self.assertFalse(e.is_contiguous(), 'E is contiguous')
        torch.linalg.eig(X, True, out=(e, v))
        Xhat = torch.mm(torch.mm(v, torch.diag(e.select(1, 0))), v.t())
        self.assertEqual(X, Xhat, 1e-8, 'VeV\' wrong')
//  [python]
def pca_torch(X, no_dims=50):
    print("Preprocessing the data using PCA...")
    (n, d) = X.shape
    X = X - torch.mean(X, 0)

    (l, M) = torch.eig(torch.mm(X.t(), X), True)
    # split M real
    # this part may be some difference for complex eigenvalue
    # but complex eignevalue is meanless here, so they are replaced by their real part
    i = 0
    while i < d:
        if l[i, 1] != 0:
            M[:, i + 1] = M[:, i]
            i += 2
        else:
            i += 1

    Y = torch.mm(X, M[:, 0:no_dims])
    return Y
//  [python]
def pca_torch(X, no_dims=50):
    print("Preprocessing the data using PCA...")
    (n, d) = X.shape
    X = X - torch.mean(X, 0)

    (l, M) = torch.linalg.eig(torch.mm(X.t(), X))
    # split M real
    # this part may be some difference for complex eigenvalue
    # but complex eignevalue is meanless here, so they are replaced by their real part
    i = 0
    while i < d:
        if l[i, 1] != 0:
            M[:, i + 1] = M[:, i]
            i += 2
        else:
            i += 1

    Y = torch.mm(X, M[:, 0:no_dims])
    return Y
//  [python]
def logdet(x):
    """
    Args:
        x: 2D positive semidefinite matrix.
    Returns: log determinant of x
    """
    # TODO for pytorch 2.0.4, use inside potrf for variable.
    print(torch.log(torch.eig(x.data)[0]))
    print(x)
    u_chol = x.potrf()
    return torch.sum(torch.log(u_chol.diag())) * 2
//  [python]
def logdet(x):
    """
    Args:
        x: 2D positive semidefinite matrix.
    Returns: log determinant of x
    """
    # TODO for pytorch 2.0.4, use inside potrf for variable.
    print(torch.log(torch.linalg.eig(x.data)[0]))
    print(x)
    u_chol = x.potrf()
    return torch.sum(torch.log(u_chol.diag())) * 2
//  [python]
def test_old_eig_out(self, device, dtype):
        # the out version of torch.eig needs to be tested manually: we can't
        # use the "test_out=True" parameter to tensor_op_tests because the
        # signature is irregular (since we have *two* output vectors)
        t = torch.randn(10, 10, dtype=dtype, device=device)
        evals, evecs = torch.eig(t, eigenvectors=True)
        #
        # check that the out= version computes the same values as the normal one
        out_evals = torch.empty_like(evals)
        out_evecs = torch.empty_like(evecs)
        evals2, evecs2 = torch.eig(t, eigenvectors=True, out=(out_evals, out_evecs))
        # check that the out tensors were used in-place
        self.assertEqual(evals2.data_ptr(), out_evals.data_ptr())
        self.assertEqual(evecs2.data_ptr(), out_evecs.data_ptr())
        # check that the result is the same as the non-out version
        self.assertEqual(evals, out_evals)
        self.assertEqual(evecs, out_evecs)
        #
        # check what happens in the eigenvectors=False case
        out_evals = torch.empty_like(evals)
        out_evecs = torch.tensor([1, 2, 3], dtype=dtype, device=device)
        evals2, evecs2 = torch.eig(t, eigenvectors=False, out=(out_evals, out_evecs))
        # check that the out_evals was used in-place
        self.assertEqual(evals2.data_ptr(), out_evals.data_ptr())
        self.assertEqual(evals, out_evals)
        # check that out_evecs was NOT touched at all
        assert out_evecs.tolist() == [1, 2, 3]
        #
        # check that we complain if we pass an out vector of the wrong dtype
        wrong_out = torch.empty((0, 0), dtype=int)
        with self.assertRaisesRegex(RuntimeError, r"Expected .* but got .*"):
            torch.eig(t, eigenvectors=True, out=(wrong_out, out_evecs))
        with self.assertRaisesRegex(RuntimeError, r"Expected .* but got .*"):
            torch.eig(t, eigenvectors=True, out=(out_evals, wrong_out))
//  [python]
def test_old_eig_out(self, device, dtype):
        # the out version of torch.linalg.eig needs to be tested manually: we can't
        # use the "test_out=True" parameter to tensor_op_tests because the
        # signature is irregular (since we have *two* output vectors)
        t = torch.randn(10, 10, dtype=dtype, device=device)
        evals, evecs = torch.linalg.eig(t)
        #
        # check that the out= version computes the same values as the normal one
        out_evals = torch.empty_like(evals)
        out_evecs = torch.empty_like(evecs)
        evals2, evecs2 = torch.linalg.eig(t, out=(out_evals, out_evecs))
        # check that the out tensors were used in-place
        self.assertEqual(evals2.data_ptr(), out_evals.data_ptr())
        self.assertEqual(evecs2.data_ptr(), out_evecs.data_ptr())
        # check that the result is the same as the non-out version
        self.assertEqual(evals, out_evals)
        self.assertEqual(evecs, out_evecs)
        #
        # check what happens in the eigenvectors=False case
        out_evals = torch.empty_like(evals)
        out_evecs = torch.tensor([1, 2, 3], dtype=dtype, device=device)
        evals2, evecs2 = torch.linalg.eig(t, out=(out_evals, out_evecs))
        # check that the out_evals was used in-place
        self.assertEqual(evals2.data_ptr(), out_evals.data_ptr())
        self.assertEqual(evals, out_evals)
        # check that out_evecs was NOT touched at all
        assert out_evecs.tolist() == [1, 2, 3]
        #
        # check that we complain if we pass an out vector of the wrong dtype
        wrong_out = torch.empty((0, 0), dtype=int)
        with self.assertRaisesRegex(RuntimeError, r"Expected .* but got .*"):
            torch.linalg.eig(t, out=(wrong_out, out_evecs))
        with self.assertRaisesRegex(RuntimeError, r"Expected .* but got .*"):
            torch.linalg.eig(t, out=(out_evals, wrong_out))
//  [python]
def isPSD(A, tol=1e-5):
    # E,V = eigh(A)
    E, V = torch.eig(A, eigenvectors = False)
    E = E[:, 0].squeeze()
    # print('min_eig = ', np.min(E) , 'max_eig = ', np.max(E), ' min_diag = ', np.min(np.diag(A)))
    # make sure symmetric positive definite
    return torch.all(E > -tol) & torch.allclose(A, A.T, atol = tol), torch.min(E)
//  [python]
def isPSD(A, tol=1e-5):
    E, V = torch.linalg.eig(A)
    E = E[:, 0].squeeze()
    return torch.all(E > -tol) & torch.allclose(A, A.T, atol = tol), torch.min(E)
//  [python]
def infer_dimension(self, update_max_rank=True, use_delta=True):
        if use_delta:
            delta = self.subspace.delta

        _, var, subspace = self.get_space()
        subspace /= (self.n_models.item() - 1) ** 0.5
        tr_sigma = var.sum()

        spectrum, _ = torch.eig(subspace @ subspace.t())
        spectrum, _ = torch.sort(spectrum[:,0], descending=True)
        
        new_max_rank, ll, _ = _infer_dimension_(spectrum.numpy(), tr_sigma.numpy(),
                                                self.n_models.item(), self.num_parameters, delta)

        if new_max_rank + 1 == self.subspace.max_rank and update_max_rank:
            self.subspace.max_rank += 1
//  [python]
def infer_dimension(self, update_max_rank=True, use_delta=True):
        if use_delta:
            delta = self.subspace.delta

        _, var, subspace = self.get_space()
        subspace /= (self.n_models.item() - 1) ** 0.5
        tr_sigma = var.sum()

        spectrum, _ = torch.linalg.eig(subspace @ subspace.t())
        spectrum, _ = torch.sort(spectrum[:,0], descending=True)
        
        new_max_rank, ll, _ = _infer_dimension_(spectrum.numpy(), tr_sigma.numpy(),
                                                self.n_models.item(), self.num_parameters, delta)

        if new_max_rank + 1 == self.subspace.max_rank and update_max_rank:
            self.subspace.max_rank += 1
//  [python]
def find_clostest_PSD(A):
    """\
    https://math.stackexchange.com/questions/648809/how-to-find-closest-positive-definite-matrix-of-non-symmetric-matrix
    Positive definite matrices are a open set (Positive semidefinite is closed), no cloest point, uses 1e-4 for the minimum eigenvalue
    """
    # cannot make sure that U = V.t() with svd,
    # U, S, V = torch.svd(A, some=False)
    S, U = torch.eig(A, eigenvectors = True)
    S = S[:, 0]
    S[S <= 0] = 1e-3
    return U @ torch.diag(S) @ U.t()
//  [python]
def find_clostest_PSD(A):
    """\
    https://math.stackexchange.com/questions/648809/how-to-find-closest-positive-definite-matrix-of-non-symmetric-matrix
    Positive definite matrices are a open set (Positive semidefinite is closed), no cloest point, uses 1e-4 for the minimum eigenvalue
    """
    # cannot make sure that U = V.t() with svd,
    # U, S, V = torch.svd(A, some=False)
    S, U = torch.linalg.eig(A)
    S = S.real
    S[S <= 0] = 1e-3
    return U @ torch.diag(S) @ U.T
//  [python]
def pca(self, X, n_components=3):
        n = X.shape[0]
        mean = torch.mean(X, dim=0)
        X = X - mean
        covariance_matrix = (1 / n) * torch.matmul(X.T, X).float()  # An old torch bug: matmul float32->float16, 
        eigenvalues, eigenvectors = torch.eig(covariance_matrix, eigenvectors=True)
        eigenvalues = torch.norm(eigenvalues, dim=1)
        idx = torch.argsort(-eigenvalues)
        eigenvectors = eigenvectors[:, idx]
        proj_mat = eigenvectors[:, 0:n_components]
        
        return proj_mat
//  [python]
def pca(self, X, n_components=3):
        n = X.shape[0]
        mean = torch.mean(X, dim=0)
        X = X - mean
        covariance_matrix = (1 / n) * torch.matmul(X.T, X).float()  # An old torch bug: matmul float32->float16, 
        eigenvalues, eigenvectors = torch.linalg.eig(covariance_matrix)
        eigenvalues = torch.norm(eigenvalues, dim=1)
        idx = torch.argsort(-eigenvalues)
        eigenvectors = eigenvectors[:, idx]
        proj_mat = eigenvectors[:, 0:n_components]
        
        return proj_mat
//  [python]
def compute_corner_impulses(self, sp, i=-1) -> Tensor:
        interaction = self.interactions[0]
        bases_n = self.G_bases.shape[0]

        batch_n = interaction.batch_n()
        k = interaction.contact_n()

        lcp_mat, lcp_vec = self.compute_lcp(sp, i=i)

        Id = torch.eye(lcp_mat.shape[1])
        Id = Id.unsqueeze(0).repeat(batch_n, 1, 1)
        lcp_mat = lcp_mat + Id * 0.001

        for mat in lcp_mat:
            real_eigs = torch.eig(mat).eigenvalues[:, 0]
            assert(torch.all(real_eigs > -torch.ones(real_eigs.shape) * 0.01))

        sols = solver_utils.solve_lcp_batch(lcp_mat, lcp_vec)

        lambda_n = sols[:, 0:k]
        lambda_t = sols[:, k:k + k * bases_n]

        lambda_t = self.compute_G().bmm(lambda_t)

        return torch.cat((lambda_n, lambda_t), dim=1)
//  [python]
def compute_corner_impulses(self, sp, i=-1) -> Tensor:
        interaction = self.interactions[0]
        bases_n = self.G_bases.shape[0]

        batch_n = interaction.batch_n()
        k = interaction.contact_n()

        lcp_mat, lcp_vec = self.compute_lcp(sp, i=i)

        Id = torch.eye(lcp_mat.shape[1])
        Id = Id.unsqueeze(0).repeat(batch_n, 1, 1)
        lcp_mat = lcp_mat + Id * 0.001

        for mat in lcp_mat:
            real_eigs = torch.linalg.eig(mat).eigenvalues[:, 0]
            assert(torch.all(real_eigs > -torch.ones(real_eigs.shape) * 0.01))

        sols = solver_utils.solve_lcp_batch(lcp_mat, lcp_vec)

        lambda_n = sols[:, 0:k]
        lambda_t = sols[:, k:k + k * bases_n]

        lambda_t = self.compute_G().bmm(lambda_t)

        return torch.cat((lambda_n, lambda_t), dim=1)
//  [python]
def pca(data, k=2):
    """
    Performs principal component analysis on the data. More
    specifically, calculates the k principal components of the data
    matrix and returns the projected data along with the eigen values
    and vectors.

    data: ndarray or torch tensor (T,N)
        the data should be organized such that T is the number
        of datapoints and N is the number of features. In
        neuro, this would mean T is time or trial and N is
        neuron.
    k: int
        number of principal components to calculate. In reality,
        all components are calculated regardless of this argument,
        so reducing k does not reduce compute time.
    """
    cent = data-data.mean(0)
    cov = cent.T@cent
    if isinstance(data, np.ndarray):
        vals, vecs = np.linalg.eig(cov)
        args = np.argsort(-vals)
    else:
        vals, vecs = torch.eig(cov,eigenvectors=True)
        vecs = vecs.T # torch returns left eigenvectors, we want right
        vals = vals[:,0] # Only reals
        args = torch.argsort(vals)
    k = min(len(args),k)
    vecs = vecs[:,args[:k]]
    vals = vals[args[:k]]
    return cent@vecs, vals, vecs
//  [python]
def pca(data, k=2):
    """
    Performs principal component analysis on the data. More
    specifically, calculates the k principal components of the data
    matrix and returns the projected data along with the eigen values
    and vectors.

    data: ndarray or torch tensor (T,N)
        the data should be organized such that T is the number
        of datapoints and N is the number of features. In
        neuro, this would mean T is time or trial and N is
        neuron.
    k: int
        number of principal components to calculate. In reality,
        all components are calculated regardless of this argument,
        so reducing k does not reduce compute time.
    """
    cent = data-data.mean(0)
    cov = cent.T@cent
    if isinstance(data, np.ndarray):
        vals, vecs = np.linalg.eig(cov)
        args = np.argsort(-vals)
    else:
        vals, vecs = torch.linalg.eig(cov)
        vecs = vecs.T # torch returns left eigenvectors, we want right
        vals = vals.real # Only reals
        args = torch.argsort(vals,descending=True)
    k = min(len(args),k)
    vecs = vecs[:,args[:k]]
    vals = vals[args[:k]]
    return cent@vecs, vals, vecs
//  [python]
def compute_impulse_lcp(self, lcp_mat, lcp_vec):
        lcp_mat = lcp_mat
        lcp_vec = lcp_vec
        Id = torch.eye(lcp_mat.shape[1])
        Id = Id.unsqueeze(0).repeat(1, 1, 1)
        lcp_mat = lcp_mat + Id * 0.001

        for mat in lcp_mat:
            real_eigs = torch.eig(mat).eigenvalues[:, 0]
            assert(torch.all(real_eigs > -torch.ones(real_eigs.shape) * 0.01))

        return solver_utils.solve_lcp_batch(lcp_mat, lcp_vec)
//  [python]
def compute_impulse_lcp(self, lcp_mat, lcp_vec):
        lcp_mat = lcp_mat
        lcp_vec = lcp_vec
        Id = torch.eye(lcp_mat.shape[1])
        Id = Id.unsqueeze(0).repeat(1, 1, 1)
        lcp_mat = lcp_mat + Id * 0.001

        for mat in lcp_mat:
            real_eigs = torch.linalg.eig(mat).eigenvalues[:, 0]
            assert(torch.all(real_eigs > -torch.ones(real_eigs.shape) * 0.01))

        return solver_utils.solve_lcp_batch(lcp_mat, lcp_vec)
//  [python]
def post_processing_spd_cholesky_torch(x_chol, min_eigenvalue, max_eigenvalue):
    """
    This function post-processes a symmetric matrix, so that its eigenvalues lie in the defined bounds.

    Parameters
    ----------
    :param x_vec: symmetric matrix in Mandel vector form
    :param min_eigenvalue: minimum eigenvalue to satisfy the constraint
    :param max_eigenvalue: maximum eigenvalue to satisfy the constraint

    Returns
    -------
    :return: symmetric matrix in Mandel vector form
    """
    # Initial shape
    init_shape = list(x_chol.shape)
    x_chol = x_chol.view(-1, init_shape[-1])

    # Dimension of SPD matrix
    dim_vec = x_chol.shape[-1]
    dim = int((-1.0 + (1.0 + 8.0 * dim_vec) ** 0.5) / 2.0)
    # Indices for Cholesky decomposition
    indices = np.tril_indices(dim)

    for n in range(x_chol.shape[0]):
        # SPD matrix
        xL = torch.zeros((dim, dim), dtype=x_chol.dtype)
        xL[indices] = x_chol[n]
        x_mat = torch.mm(xL, xL.T)

        # Check constraints
        eigdec = torch.eig(x_mat, eigenvectors=True)
        eigvals = eigdec.eigenvalues[:, 0]
        eigvecs = eigdec.eigenvectors

        eigvals[eigvals <= min_eigenvalue] = min_eigenvalue  # PD constraint
        eigvals[eigvals > max_eigenvalue] = max_eigenvalue  # Max eigenvalue constraint

        x_mat = torch.mm(torch.mm(eigvecs, torch.diag(eigvals)), torch.inverse(eigvecs))

        # Cholesky decomposition
        x_chol[n] = torch.cholesky(x_mat)[indices]

    return x_chol.view(init_shape)
//  [python]
def post_processing_spd_cholesky_torch(x_chol, min_eigenvalue, max_eigenvalue):
    """
    This function post-processes a symmetric matrix, so that its eigenvalues lie in the defined bounds.

    Parameters
    ----------
    :param x_vec: symmetric matrix in Mandel vector form
    :param min_eigenvalue: minimum eigenvalue to satisfy the constraint
    :param max_eigenvalue: maximum eigenvalue to satisfy the constraint

    Returns
    -------
    :return: symmetric matrix in Mandel vector form
    """
    # Initial shape
    init_shape = list(x_chol.shape)
    x_chol = x_chol.view(-1, init_shape[-1])

    # Dimension of SPD matrix
    dim_vec = x_chol.shape[-1]
    dim = int((-1.0 + (1.0 + 8.0 * dim_vec) ** 0.5) / 2.0)
    # Indices for Cholesky decomposition
    indices = np.tril_indices(dim)

    for n in range(x_chol.shape[0]):
        # SPD matrix
        xL = torch.zeros((dim, dim), dtype=x_chol.dtype)
        xL[indices] = x_chol[n]
        x_mat = torch.mm(xL, xL.T)

        # Check constraints
        eigdec = torch.linalg.eig(x_mat)
        eigvals = eigdec.eigenvalues[:, 0]
        eigvecs = eigdec.eigenvectors

        eigvals[eigvals <= min_eigenvalue] = min_eigenvalue  # PD constraint
        eigvals[eigvals > max_eigenvalue] = max_eigenvalue  # Max eigenvalue constraint

        x_mat = torch.mm(torch.mm(eigvecs, torch.diag(eigvals)), torch.inverse(eigvecs))

        # Cholesky decomposition
        x_chol[n] = torch.cholesky(x_mat)[indices]

    return x_chol.view(init_shape)
//  [python]
def test_eig_no_eigenvectors(self):
        A = torch.tensor([[1., 2.], [2., 4.]], dtype=torch.float32, requires_grad=True)
        w, v = torch.eig(A, eigenvectors=False)
        with self.assertRaisesRegex(RuntimeError, 'is not differentiable'):
            torch.autograd.backward([w, v], [torch.ones_like(w), torch.ones_like(v)])
//  [python]
def test_eig_no_eigenvectors(self):
        A = torch.tensor([[1., 2.], [2., 4.]], dtype=torch.float32, requires_grad=True)
        w, v = torch.linalg.eig(A)
        with self.assertRaisesRegex(RuntimeError, 'is not differentiable'):
            torch.autograd.backward([w, v], [torch.ones_like(w), torch.ones_like(v)])
//  [python]
def test_eig_complex_eigenvalues(self):
        A = torch.tensor([[0., -1.], [1., 0.]], dtype=torch.float32, requires_grad=True)
        w, v = torch.eig(A, eigenvectors=True)
        with self.assertRaisesRegex(RuntimeError, 'does not support complex eigenvalues'):
            torch.autograd.backward([w, v], [torch.ones_like(w), torch.ones_like(v)])
//  [python]
def test_eig_complex_eigenvalues(self):
        A = torch.tensor([[0., -1.], [1., 0.]], dtype=torch.float32, requires_grad=True)
        w, v = torch.linalg.eig(A)
        with self.assertRaisesRegex(RuntimeError, 'does not support complex eigenvalues'):
            torch.autograd.backward([w, v], [torch.ones_like(w), torch.ones_like(v)])
//  [python]
def modeig_forward_etc(P, op, eig_mode='svd', param=None):
    '''
    Generic forward function of non-linear eigenvalue modification
    LogEig, ReEig, etc inherit from this class
    Input P: (batch_size,channels) SPD matrices of size (n,n)
    Output X: (batch_size,channels) modified symmetric matrices of size (n,n)
    '''
    batch_size, channels, n, n = P.shape
    U, S = torch.zeros_like(P, device=P.device), torch.zeros(batch_size,
                                                             channels,
                                                             n,
                                                             dtype=P.dtype,
                                                             device=P.device)
    for i in range(batch_size):
        for j in range(channels):
            if (eig_mode == 'eig'):
                # s, U[i, j] = torch.linalg.eig(P[i, j])
                s, U[i, j] = torch.eig(P[i, j].cpu(), True)
                S[i, j] = s[:, 0]
            elif (eig_mode == 'svd'):
                # U[i, j], S[i, j], _ = torch.svd(add_id_matrix(P[i, j], 1e-5))
                U[i, j], S[i, j], _ = torch.svd(add_id_matrix(P[i, j].cpu(), 1e-5))
    U, S = U.cuda(), S.cuda()
    S_fn = op.fn(S, param)
    X = U.matmul(BatchDiag(S_fn)).matmul(U.transpose(2, 3))
    return X, U, S, S_fn
//  [python]
def modeig_forward_etc(P, op, eig_mode='svd', param=None):
    '''
    Generic forward function of non-linear eigenvalue modification
    LogEig, ReEig, etc inherit from this class
    Input P: (batch_size,channels) SPD matrices of size (n,n)
    Output X: (batch_size,channels) modified symmetric matrices of size (n,n)
    '''
    batch_size, channels, n, n = P.shape
    U, S = torch.zeros_like(P, device=P.device), torch.zeros(batch_size,
                                                             channels,
                                                             n,
                                                             dtype=P.dtype,
                                                             device=P.device)
    for i in range(batch_size):
        for j in range(channels):
            if (eig_mode == 'eig'):
                s, U[i, j] = torch.linalg.eig(P[i, j])
                S[i, j] = s[:, 0]
            elif (eig_mode == 'svd'):
                U[i, j], S[i, j], _ = torch.svd(add_id_matrix(P[i, j].cpu(), 1e-5))
    U, S = U.cuda(), S.cuda()
    S_fn = op.fn(S, param)
    X = U.matmul(BatchDiag(S_fn)).matmul(U.transpose(2, 3))
    return X, U, S, S_fn
//  [python]
def is_psd(self, mat):
        eigenvalues = torch.eig(mat)[0][:, 0].sort(descending=True)[0]
        return eigenvalues, float((mat == mat.t()).all() and (eigenvalues >= 0).all())
//  [python]
def is_psd(self, mat):
        eigenvalues = torch.linalg.eig(mat).eigenvalues.sort(descending=True)[0]
        return eigenvalues, float((mat == mat.t()).all() and (eigenvalues >= 0).all())
//  [python]
def __init__(self, smallest_eigen=True):
        super().__init__(smallest_eigen)
         
        self.solvers = {
            'torch_cuda': lambda m: torch.eig(m, eigenvectors=True),
        }
//  [python]
def __init__(self, smallest_eigen=True):
        super().__init__(smallest_eigen)
         
        self.solvers = {
            'torch_cuda': lambda m: torch.linalg.eig(m),
        }
//  [python]
def Dist_Comp(self, outputs):


        self.M = outputs
        self.k = self.M
        #self.M.requires_grad = False
        self.y = self.M
        n = self.k.size(0)
        m = self.y.size(0)
        d = self.k.size(1)

        self.k = self.k.unsqueeze(1).expand(n, m, d)
        self.y = self.y.unsqueeze(0).expand(n, m, d)
        self.dist = torch.pow(self.k - self.y, 2).sum(2)

        self.A = 1 - self.dist
        mi = torch.min(self.A)
        te = 0.001  # a parameter to set to the smallest (negative) value in the matrix
        self.A = mi + self.A + te
        # z = torch.diag(torch.tensor([float(0)]).expand(self.l)).cuda()
        # setting the diagonal to zero

        z = -1 * torch.diag(self.A)
        z = torch.diag(z.type(torch.float))
        self.A = z + self.A
        B_size = self.A.size(0)

        self.F_rank = Variable((torch.zeros(B_size, B_size).cuda()))
        #B_size = self.B_size
        for i in range(0,B_size):

            v = torch.ones(B_size,1)
            v[i] = 0
            ind = (v!=0).nonzero()


            ind = ind[:, 0]
            M=self.A[ind][:,ind]


            ########### save the mat as matlab file to cross

            EE=torch.eig(M)
            EE2 = EE[0]

            alpha = torch.max(EE2[:, 0])
            alpha = alpha + 0.99
            ident_M = torch.eye(B_size,B_size)
            ident_M[i,i] = 0
            ident_M= ident_M*float(alpha)
            self.A_2= self.A - ident_M.cuda()
            self.A_2=self.A_2+alpha
            self.A_2 = self.A_2/torch.max(self.A_2)
            # calling Replicator dynamics
            P_g_rank= self.Replicator()
            P_g_rank = P_g_rank.view(1,-1)
            self.F_rank[i,:] = P_g_rank

        return self.F_rank
//  [python]
def Dist_Comp(self, outputs):


        self.M = outputs
        self.k = self.M
        #self.M.requires_grad = False
        self.y = self.M
        n = self.k.size(0)
        m = self.y.size(0)
        d = self.k.size(1)

        self.k = self.k.unsqueeze(1).expand(n, m, d)
        self.y = self.y.unsqueeze(0).expand(n, m, d)
        self.dist = torch.pow(self.k - self.y, 2).sum(2)

        self.A = 1 - self.dist
        mi = torch.min(self.A)
        te = 0.001  # a parameter to set to the smallest (negative) value in the matrix
        self.A = mi + self.A + te
        # z = torch.diag(torch.tensor([float(0)]).expand(self.l)).cuda()
        # setting the diagonal to zero

        z = -1 * torch.diag(self.A)
        z = torch.diag(z.type(torch.float))
        self.A = z + self.A
        B_size = self.A.size(0)

        self.F_rank = Variable((torch.zeros(B_size, B_size).cuda()))
        #B_size = self.B_size
        for i in range(0,B_size):

            v = torch.ones(B_size,1)
            v[i] = 0
            ind = (v!=0).nonzero()


            ind = ind[:, 0]
            M=self.A[ind][:,ind]


            ########### save the mat as matlab file to cross

            EE=torch.linalg.eig(M)
            EE2 = EE[0]

            alpha = torch.max(EE2[:, 0])
            alpha = alpha + 0.99
            ident_M = torch.eye(B_size,B_size)
            ident_M[i,i] = 0
            ident_M= ident_M*float(alpha)
            self.A_2= self.A - ident_M.cuda()
            self.A_2=self.A_2+alpha
            self.A_2 = self.A_2/torch.max(self.A_2)
            # calling Replicator dynamics
            P_g_rank= self.Replicator()
            P_g_rank = P_g_rank.view(1,-1)
            self.F_rank[i,:] = P_g_rank

        return self.F_rank
//  [python]
def get_condition_number(matrix):
    """Calculate condition number of a matrix (to check if its ill-conditioned)"""
    values, _ = torch.eig(matrix)
    real_values = values[:, 0]
    max_eig_value = torch.max(torch.abs(real_values))
    min_eig_value = torch.min(torch.abs(real_values))
    cond_num = (max_eig_value / min_eig_value).item()
    return cond_num
//  [python]
def get_condition_number(matrix):
    """Calculate condition number of a matrix (to check if its ill-conditioned)"""
    values, _ = torch.linalg.eig(matrix)
    real_values = values[:, 0]
    max_eig_value = torch.max(torch.abs(real_values))
    min_eig_value = torch.min(torch.abs(real_values))
    cond_num = (max_eig_value / min_eig_value).item()
    return cond_num
//  [python]
def A2normed_scaled_L(A):
    """
    From original A(not A+I) to normalized, rescaled Graph Laplacian
    """
    ## L = I - D^-0.5 * A * D^-0.5
    normed_L = torch.eye(A.shape[0]) - normalize_undigraph(A)
    ## L_ = L * (2/eigval_max) - I
    # eigenvalue: fetch the real part
    lambda_max = 1.02 * torch.max(torch.eig(normed_L, eigenvectors=False)[0][:, 0])
    normed_rescaled_L = normed_L * (2/lambda_max) - torch.eye(A.shape[0])

    return normed_rescaled_L
//  [python]
def A2normed_scaled_L(A):
    """
    From original A(not A+I) to normalized, rescaled Graph Laplacian
    """
    ## L = I - D^-0.5 * A * D^-0.5
    normed_L = torch.eye(A.shape[0]) - normalize_undigraph(A)
    ## L_ = L * (2/eigval_max) - I
    # eigenvalue: fetch the real part
    lambda_max = 1.02 * torch.max(torch.linalg.eigvals(normed_L).real)
    normed_rescaled_L = normed_L * (2/lambda_max) - torch.eye(A.shape[0])

    return normed_rescaled_L
//  [python]
def rescale_laplacian(L):
        # rescale laplacian to arccos range [-1,1] for input to Chebyshev polynomials of the first kind
        try:
            lambda_ = torch.eig(L)[0][:,0]      # get the real parts of eigenvalues
            lambda_max = lambda_.max()      # get the largest eigenvalue
        except:
            print("Eigen_value calculation didn't converge, using max_eigen_val=2 instead.")
            lambda_max = 2
        L_rescaled = (2 / lambda_max) * L - torch.eye(L.shape[0])
        return L_rescaled
//  [python]
def rescale_laplacian(L):
        # rescale laplacian to arccos range [-1,1] for input to Chebyshev polynomials of the first kind
        try:
            lambda_ = torch.linalg.eig(L).eigenvalues.real      # get the real parts of eigenvalues
            lambda_max = lambda_.max()      # get the largest eigenvalue
        except:
            print("Eigen_value calculation didn't converge, using max_eigen_val=2 instead.")
            lambda_max = 2
        L_rescaled = (2 / lambda_max) * L - torch.eye(L.shape[0])
        return L_rescaled
//  [python]
def log_SO3_eigen(C):  # no autodiff
    phi_norm = torch.acos(torch.clamp((torch.trace(C) - 1) / 2, -1.0, 1.0))

    # eig is not very food for C close to identity, will only keep around 3 decimals places
    w, v = torch.eig(C, eigenvectors=True)
    a = torch.tensor([0., 0., 0.], device=C.device)
    for i in range(0, w.size(0)):
        if torch.abs(w[i, 0] - 1.0) < 1e-6 and torch.abs(w[i, 1] - 0.0) < 1e-6:
            a = v[:, i]

    assert (torch.abs(torch.norm(a) - 1.0) < 1e-6)

    if torch.allclose(exp_SO3(phi_norm * a), C, atol=1e-3):
        return phi_norm * a
    elif torch.allclose(exp_SO3(-phi_norm * a), C, atol=1e-3):
        return -phi_norm * a
    else:
        raise ValueError("Invalid logarithmic mapping")
//  [python]
def log_SO3_eigen(C):  # no autodiff
    phi_norm = torch.acos(torch.clamp((torch.trace(C) - 1) / 2, -1.0, 1.0))

    # eig is not very food for C close to identity, will only keep around 3 decimals places
    w, v = torch.linalg.eig(C)
    a = torch.tensor([0., 0., 0.], device=C.device)
    for i in range(0, w.size(0)):
        if torch.abs(w[i, 0] - 1.0) < 1e-6 and torch.abs(w[i, 1] - 0.0) < 1e-6:
            a = v[:, i]

    assert (torch.abs(torch.norm(a) - 1.0) < 1e-6)

    if torch.allclose(exp_SO3(phi_norm * a), C, atol=1e-3):
        return phi_norm * a
    elif torch.allclose(exp_SO3(-phi_norm * a), C, atol=1e-3):
        return -phi_norm * a
    else:
        raise ValueError("Invalid logarithmic mapping")
//  [python]
def test_eig(self):
        def func(B):
            return torch.eig(B, eigenvectors=True)

        def func_eigvals(B):
            return torch.eig(B, eigenvectors=True)[0]

        def func_eigvecs(B):
            return torch.eig(B, eigenvectors=True)[1]

        def run_test(dims):
            # The backward operation for eig only works for real eigenvalues,
            # so the matrix should be B = U^{-1}*A*U where A is a random
            # symmetric matrix and U is a random full-rank matrix.
            # Slight change to the matrix should not make the eigenvalues
            # complex, so we apply requires_grad_ to B, not A and U

            A = random_symmetric_matrix(dims[-1], *dims[:-2])
            U = torch.rand(*dims)
            Uinv = torch.inverse(U)
            B = torch.matmul(Uinv, torch.matmul(A, U)).requires_grad_()

            gradcheck(func, [B])
            gradgradcheck(func, [B])
            gradcheck(func_eigvals, [B])
            gradgradcheck(func_eigvals, [B])
            gradcheck(func_eigvecs, [B])
            gradgradcheck(func_eigvecs, [B])

        for dims in [(3, 3), (5, 5)]:
            run_test(dims)
//  [python]
def test_eig(self):
        def func(B):
            return torch.linalg.eig(B)

        def func_eigvals(B):
            return torch.linalg.eig(B)[0]

        def func_eigvecs(B):
            return torch.linalg.eig(B)[1]

        def run_test(dims):
            # The backward operation for eig only works for real eigenvalues,
            # so the matrix should be B = U^{-1}*A*U where A is a random
            # symmetric matrix and U is a random full-rank matrix.
            # Slight change to the matrix should not make the eigenvalues
            # complex, so we apply requires_grad_ to B, not A and U

            A = random_symmetric_matrix(dims[-1], *dims[:-2])
            U = torch.rand(*dims)
            Uinv = torch.inverse(U)
            B = torch.matmul(Uinv, torch.matmul(A, U)).requires_grad_()

            gradcheck(func, [B])
            gradgradcheck(func, [B])
            gradcheck(func_eigvals, [B])
            gradgradcheck(func_eigvals, [B])
            gradcheck(func_eigvecs, [B])
            gradgradcheck(func_eigvecs, [B])

        for dims in [(3, 3), (5, 5)]:
            run_test(dims)
//  [python]
def pca_torch(X, no_dims=50):
    print("=> preprocessing the features using PCA...")
    (n, d) = X.shape
    X = X - torch.mean(X, 0)

    (l, M) = torch.eig(torch.mm(X.t(), X), True)
    # split M real
    for i in range(d):
        if l[i, 1] != 0:
            M[:, i+1] = M[:, i]
            i += 1

    Y = torch.mm(X, M[:, 0:no_dims])
    return Y
//  [python]
def pca_torch(X, no_dims=50):
    print("=> preprocessing the features using PCA...")
    (n, d) = X.shape
    X = X - torch.mean(X, 0)

    (l, M) = torch.linalg.eig(torch.mm(X.t(), X))
    # split M real
    for i in range(d):
        if l[i, 1] != 0:
            M[:, i+1] = M[:, i]
            i += 1

    Y = torch.mm(X, M[:, 0:no_dims])
    return Y
//  [python]
def pca_torch(self, X):
        covariance = torch.transpose(X, 1, 0) @ X
        S, U = torch.eig(covariance, eigenvectors=True)
        return S, U