jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def validate(self, validation_past, validation_future):

        """
        Validate the model using validation dataset to get validation score.
        """

        self._lobe.eval()
        if not self._network_type == 'SRVNet':
            self._lobe_timelagged.eval()

        if isinstance(validation_past, np.ndarray):
            validation_past = torch.from_numpy(validation_past.astype(np.float32))
        if isinstance(validation_future, np.ndarray):
            validation_future = torch.from_numpy(validation_future.astype(np.float32))
        validation_past = validation_past.to(device=self._device)
        validation_future = validation_future.to(device=self._device)

        with torch.no_grad():
            if self._network_type == "SRVNet":
                _val_pastproject = self._lobe(validation_past)
                _val_futureproject = self._lobe(validation_future)
            else:
                _val_pastproject = self._lobe(validation_past)
                _val_futureproject = self._lobe_timelagged(validation_future)
            _val_score = loss_function(pastdata=_val_pastproject, futuredata=_val_futureproject, criterion=self._network_type, 
                              epsilon=self.__epsilon, method=self.__covariance_method, device=self._device)
            if self._network_type == 'SRVNet':
                rev_propogator = reverse_propogator(pastdata=_val_pastproject, futuredata=_val_futureproject, 
                                                    epsilon=self.__epsilon, method=self.__covariance_method, device=self._device)
                _eigenval, _ = torch.symeig(rev_propogator)
                _eigenval = torch.flip(_eigenval[-int(_val_pastproject.shape[1]):], dims=[0]).detach()
                return -_val_score.item(), _eigenval.cpu().numpy()
            else:
                return -_val_score.item()
//  [python]
def validate(self, validation_past, validation_future):

        """
        Validate the model using validation dataset to get validation score.
        """

        self._lobe.eval()
        if not self._network_type == 'SRVNet':
            self._lobe_timelagged.eval()

        if isinstance(validation_past, np.ndarray):
            validation_past = torch.from_numpy(validation_past.astype(np.float32))
        if isinstance(validation_future, np.ndarray):
            validation_future = torch.from_numpy(validation_future.astype(np.float32))
        validation_past = validation_past.to(device=self._device)
        validation_future = validation_future.to(device=self._device)

        with torch.no_grad():
            if self._network_type == "SRVNet":
                _val_pastproject = self._lobe(validation_past)
                _val_futureproject = self._lobe(validation_future)
            else:
                _val_pastproject = self._lobe(validation_past)
                _val_futureproject = self._lobe_timelagged(validation_future)
            _val_score = loss_function(pastdata=_val_pastproject, futuredata=_val_futureproject, criterion=self._network_type, 
                              epsilon=self.__epsilon, method=self.__covariance_method, device=self._device)
            if self._network_type == 'SRVNet':
                rev_propogator = reverse_propogator(pastdata=_val_pastproject, futuredata=_val_futureproject, 
                                                    epsilon=self.__epsilon, method=self.__covariance_method, device=self._device)
                _eigenval, _ = torch.linalg.eigh(rev_propogator)
                _eigenval = torch.flip(_eigenval[-int(_val_pastproject.shape[1]):], dims=[0]).detach()
                return -_val_score.item(), _eigenval.cpu().numpy()
            else:
                return -_val_score.item()
//  [python]
def test_eigh_lower_uplo(self, device, dtype):
        def run_test(shape, batch, uplo):
            # check lower case uplo
            # use non-symmetric input to check whether uplo argument is working as intended
            matrix = torch.randn(shape, shape, *batch, dtype=dtype, device=device)
            expected_w, expected_v = np.linalg.eigh(matrix.cpu().numpy(), UPLO=uplo)
            actual_w, actual_v = torch.symeig(matrix, eigenvectors=True, upper=uplo == "u")
            self.assertEqual(actual_w, expected_w)
            self.assertEqual(abs(actual_v), abs(expected_v))

        uplos = ["u", "l"]
        for uplo in uplos:
            run_test(3, (2, 2), uplo)
//  [python]
def test_eigh_lower_uplo(self, device, dtype):
        def run_test(shape, batch, uplo):
            # check lower case uplo
            # use non-symmetric input to check whether uplo argument is working as intended
            matrix = torch.randn(shape, shape, *batch, dtype=dtype, device=device)
            expected_w, expected_v = np.linalg.eigh(matrix.cpu().numpy(), UPLO=uplo)
            actual_w, actual_v = torch.linalg.eigh(matrix, UPLO=uplo)
            self.assertEqual(actual_w, expected_w)
            self.assertEqual(abs(actual_v), abs(expected_v))

        uplos = ["u", "l"]
        for uplo in uplos:
            run_test(3, (2, 2), uplo)
//  [python]
def loss_function(pastdata=None, futuredata=None, criterion='SRVNet', epsilon=1e-6, method='cutoff', device=None):

    """
    Loss function calculation for the VAMPNet / SRVNet training.

    Parameters
    ----------
    pastdata: torch.Tensor, default=None
        The past part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    futuredata: torch.Tensor, default=None
        The future part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    criterion: string, default='SRVNet'
        The criterion used to calculate the score/loss fucntion, can be chosen from: 'SRVNet', 'VAMP1', 'VAMP2', 'VAMPE'
    epsilon: float, default=1e-6
        The epsilon modification used to calculate the inverse of the instantaneous matrix.
    method: string, default='cutoff'
        The method used to calculate the inverse of the instantaneous matix: : 'regularize', 'cutoff', 'clamp'
    
    Returns
    ----------
    loss function: float
        The loss function for the set criterion.
    """

    _criterion = ['SRVNet', 'VAMP1', 'VAMP2', 'VAMPE']
    assert criterion in _criterion, "Invalid criterion "+str(criterion)+" was adopted, should use criterion in: "+str(_criterion)
    
    if criterion == 'SRVNet':
        rev_propogator = reverse_propogator(pastdata=pastdata, futuredata=futuredata, epsilon=0, method='regularize', device=device)
        _eigenval, _ = torch.symeig(rev_propogator)
        _score = torch.sum(torch.pow(_eigenval, 2)) + 1


    elif criterion == 'VAMP1':
        koopman_mat = koopman_matrix(pastdata=pastdata, futuredata=futuredata, epsilon=epsilon, method=method, device=device)
        _score = torch.norm(koopman_mat, p='nuc') + 1


    elif criterion == 'VAMP2':
        koopman_mat = koopman_matrix(pastdata=pastdata, futuredata=futuredata, epsilon=epsilon, method=method, device=device)
        _score = torch.pow(torch.norm(koopman_mat, p='fro'), 2) + 1


    elif criterion == 'VAMPE':
        c00, c0t, ctt = covariance(pastdata=pastdata, futuredata=futuredata, remove_mean=True, device=device)
        c00_sqrt_inv = matrix_inverse(matrix=c00, sqrt_root=True, epsilon=epsilon, method=method)
        ctt_sqrt_inv = matrix_inverse(matrix=ctt, sqrt_root=True,  epsilon=epsilon, method=method)

        koopman_mat = torch.matmul(c00_sqrt_inv, torch.matmul(c0t, ctt_sqrt_inv)).t()
        u, sigma, v = torch.svd(koopman_mat)
        mask = sigma > epsilon
        
        u = torch.mm(c00_sqrt_inv, u[:, mask])
        v = torch.mm(ctt_sqrt_inv, v[:, mask])
        sigma = sigma[mask]

        u_t = u.t(); v_t = v.t()
        sigma = torch.diag(sigma)
        _score = torch.trace(2. * torch.linalg.multi_dot([sigma, u_t, c0t, v]) - torch.linalg.multi_dot([sigma, u_t, c00, u, sigma, v_t, ctt, v])) + 1

    return -1. * _score
//  [python]
def loss_function(pastdata=None, futuredata=None, criterion='SRVNet', epsilon=1e-6, method='cutoff', device=None):

    """
    Loss function calculation for the VAMPNet / SRVNet training.

    Parameters
    ----------
    pastdata: torch.Tensor, default=None
        The past part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    futuredata: torch.Tensor, default=None
        The future part of the time-lagged sequence data. Can be generated from TimeLaggedDataset.
    criterion: string, default='SRVNet'
        The criterion used to calculate the score/loss fucntion, can be chosen from: 'SRVNet', 'VAMP1', 'VAMP2', 'VAMPE'
    epsilon: float, default=1e-6
        The epsilon modification used to calculate the inverse of the instantaneous matrix.
    method: string, default='cutoff'
        The method used to calculate the inverse of the instantaneous matix: : 'regularize', 'cutoff', 'clamp'
    
    Returns
    ----------
    loss function: float
        The loss function for the set criterion.
    """

    _criterion = ['SRVNet', 'VAMP1', 'VAMP2', 'VAMPE']
    assert criterion in _criterion, "Invalid criterion "+str(criterion)+" was adopted, should use criterion in: "+str(_criterion)
    
    if criterion == 'SRVNet':
        rev_propogator = reverse_propogator(pastdata=pastdata, futuredata=futuredata, epsilon=0, method='regularize', device=device)
        _eigenval, _ = torch.linalg.eigh(rev_propogator)
        _score = torch.sum(torch.pow(_eigenval, 2)) + 1


    elif criterion == 'VAMP1':
        koopman_mat = koopman_matrix(pastdata=pastdata, futuredata=futuredata, epsilon=epsilon, method=method, device=device)
        _score = torch.norm(koopman_mat, p='nuc') + 1


    elif criterion == 'VAMP2':
        koopman_mat = koopman_matrix(pastdata=pastdata, futuredata=futuredata, epsilon=epsilon, method=method, device=device)
        _score = torch.pow(torch.norm(koopman_mat, p='fro'), 2) + 1


    elif criterion == 'VAMPE':
        c00, c0t, ctt = covariance(pastdata=pastdata, futuredata=futuredata, remove_mean=True, device=device)
        c00_sqrt_inv = matrix_inverse(matrix=c00, sqrt_root=True, epsilon=epsilon, method=method)
        ctt_sqrt_inv = matrix_inverse(matrix=ctt, sqrt_root=True,  epsilon=epsilon, method=method)

        koopman_mat = torch.matmul(c00_sqrt_inv, torch.matmul(c0t, ctt_sqrt_inv)).t()
        u, sigma, v = torch.svd(koopman_mat)
        mask = sigma > epsilon
        
        u = torch.mm(c00_sqrt_inv, u[:, mask])
        v = torch.mm(ctt_sqrt_inv, v[:, mask])
        sigma = sigma[mask]

        u_t = u.t(); v_t = v.t()
        sigma = torch.diag(sigma)
        _score = torch.trace(2. * torch.linalg.multi_dot([sigma, u_t, c0t, v]) - torch.linalg.multi_dot([sigma, u_t, c00, u, sigma, v_t, ctt, v])) + 1

    return -1. * _score
//  [python]
def __expm__(self, matrix, symmetric):
        r"""Calculates matrix exponential.

        Args:
            matrix (Tensor): Matrix to take exponential of.
            symmetric (bool): Specifies whether the matrix is symmetric.

        :rtype: (:class:`Tensor`)
        """
        if symmetric:
            e, V = torch.eig(matrix, eigenvectors=True)
            e = e[:,0]
            diff_mat = V @ torch.diag(e.exp()) @ V.t()
        else:
            diff_mat_np = expm(matrix.cpu().numpy())
            diff_mat = torch.Tensor(diff_mat_np).to(matrix.device)
        return diff_mat
//  [python]
def __expm__(self, matrix, symmetric):
        r"""Calculates matrix exponential.

        Args:
            matrix (Tensor): Matrix to take exponential of.
            symmetric (bool): Specifies whether the matrix is symmetric.

        :rtype: (:class:`Tensor`)
        """
        if symmetric:
            e, V = torch.linalg.eigh(matrix, UPLO='U')
            diff_mat = V @ torch.diag(e.exp()) @ V.t()
        else:
            diff_mat_np = expm(matrix.cpu().numpy())
            diff_mat = torch.Tensor(diff_mat_np).to(matrix.device)
        return diff_mat
//  [python]
def test_eigh(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_matrix

        def run_test(shape, batch, uplo):
            matrix = random_hermitian_matrix(shape, *batch, dtype=dtype, device=device)
            expected_w, expected_v = np.linalg.eigh(matrix.cpu().numpy(), UPLO=uplo)
            actual_w, actual_v = torch.linalg.eigh(matrix, UPLO=uplo)
            self.assertEqual(actual_w, expected_w)
            # sign of eigenvectors is not unique and therefore absolute values are compared
            self.assertEqual(abs(actual_v), abs(expected_v))
            # additionally we can flip the sign and then compare the values
            # let's choose the convention that the first element of the eigenvector should be positive,
            # otherwise flip the sign of the eigenvector
            if matrix.numel() > 0:
                sign = np.sign(expected_v[..., 0, :]).reshape(batch + (1, shape))
                expected_v = sign * expected_v
                torch_real_slice = actual_v[..., 0, :].real if dtype.is_complex else actual_v[..., 0, :]
                sign = torch.sign(torch_real_slice).reshape(batch + (1, shape))
                actual_v = sign * actual_v
                self.assertEqual(actual_v, expected_v)

            # check the out= variant
            out_w = torch.empty_like(actual_w)
            out_v = torch.empty_like(actual_v)
            ans_w, ans_v = torch.linalg.eigh(matrix, UPLO=uplo, out=(out_w, out_v))
            self.assertEqual(ans_w, out_w)
            self.assertEqual(ans_v, out_v)
            self.assertEqual(ans_w, actual_w)
            self.assertEqual(abs(ans_v), abs(actual_v))

        shapes = (0, 3, 5)
        batches = ((), (3, ), (2, 2))
        uplos = ["U", "L"]
        for shape, batch, uplo in itertools.product(shapes, batches, uplos):
            run_test(shape, batch, uplo)
//  [python]
def warmup(x):
    n = x.shape[0]
        
    x = F.normalize(x, dim=1) / np.sqrt(n)
    X = x @ x.T 

    eigenvalues, _ = torch.symeig(X)

    uniform = torch.ones_like(eigenvalues).to(device) / n
    loss = F.mse_loss(eigenvalues, uniform)

    return loss
//  [python]
def warmup(x):
    n = x.shape[0]
        
    x = F.normalize(x, dim=1) / np.sqrt(n)
    X = x @ x.T 

    eigenvalues, _ = torch.linalg.eigh(X)

    uniform = torch.ones_like(eigenvalues).to(device) / n
    loss = F.mse_loss(eigenvalues, uniform)

    return loss
//  [python]
def test_eigh_errors_and_warnings(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_matrix

        # eigh requires a square matrix
        t = torch.randn(2, 3, device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, "must be batches of square matrices"):
            t.eigh()

        # eigh requires 'uplo' parameter to be 'U' or 'L'
        t = torch.randn(3, 3, device=device, dtype=dtype)
        for uplo in ["a", "wrong"]:
            with self.assertRaisesRegex(RuntimeError, "be \'L\' or \'U\'"):
                t.eigh(UPLO=uplo)
            with self.assertRaisesRegex(ValueError, "be \'L\' or \'U\'"):
                np.linalg.eigh(t.cpu().numpy(), UPLO=uplo)

        # if non-empty out tensor with wrong shape is passed a warning is given
        a = random_hermitian_matrix(3, dtype=dtype, device=device)
        real_dtype = a.real.dtype if dtype.is_complex else dtype
        out_w = torch.empty(7, 7, dtype=real_dtype, device=device)
        out_v = torch.empty(7, 7, dtype=dtype, device=device)
        with warnings.catch_warnings(record=True) as w:
            # Trigger warning
            a.eigh(out=(out_w, out_v))
            # Check warning occurs
            self.assertEqual(len(w), 2)
            self.assertTrue("An output with one or more elements was resized" in str(w[-2].message))
            self.assertTrue("An output with one or more elements was resized" in str(w[-1].message))

        # dtypes should match
        out_w = torch.empty_like(a).to(torch.int)
        out_v = torch.empty_like(a)
        with self.assertRaisesRegex(RuntimeError, "dtype Int does not match self dtype"):
            a.eigh(out=(out_w, out_v))
//  [python]
def test_eigh_errors_and_warnings(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_matrix

        # eigh requires a square matrix
        t = torch.randn(2, 3, device=device, dtype=dtype)
        with self.assertRaisesRegex(RuntimeError, "must be batches of square matrices"):
            torch.linalg.eigh(t)

        # eigh requires 'uplo' parameter to be 'U' or 'L'
        t = torch.randn(3, 3, device=device, dtype=dtype)
        for uplo in ["a", "wrong"]:
            with self.assertRaisesRegex(RuntimeError, "be \'L\' or \'U\'"):
                torch.linalg.eigh(t, UPLO=uplo)
            with self.assertRaisesRegex(ValueError, "be \'L\' or \'U\'"):
                np.linalg.eigh(t.cpu().numpy(), UPLO=uplo)

        # if non-empty out tensor with wrong shape is passed a warning is given
        a = random_hermitian_matrix(3, dtype=dtype, device=device)
        real_dtype = a.real.dtype if dtype.is_complex else dtype
        out_w = torch.empty(7, 7, dtype=real_dtype, device=device)
        out_v = torch.empty(7, 7, dtype=dtype, device=device)
        with warnings.catch_warnings(record=True) as w:
            # Trigger warning
            torch.linalg.eigh(a, out=(out_w, out_v))
            # Check warning occurs
            self.assertEqual(len(w), 2)
            self.assertTrue("An output with one or more elements was resized" in str(w[-2].message))
            self.assertTrue("An output with one or more elements was resized" in str(w[-1].message))

        # dtypes should match
        out_w = torch.empty_like(a).to(torch.int)
        out_v = torch.empty_like(a)
        with self.assertRaisesRegex(RuntimeError, "dtype Int does not match self dtype"):
            torch.linalg.eigh(a, out=(out_w, out_v))
//  [python]
def test_eigh_non_contiguous(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_matrix

        def run_test(matrix, uplo):
            self.assertFalse(matrix.is_contiguous())
            expected_w, expected_v = np.linalg.eigh(matrix.cpu().numpy(), UPLO=uplo)
            actual_w, actual_v = torch.eig(matrix, eigenvectors=True, UPLO=uplo)
            self.assertEqual(actual_w, expected_w)
            # sign of eigenvectors is not unique and therefore absolute values are compared
            self.assertEqual(abs(actual_v), abs(expected_v))

        def run_test_permuted(shape, batch, uplo):
            # check for permuted / transposed inputs
            matrix = random_hermitian_matrix(shape, *batch, dtype=dtype, device=device)
            matrix = matrix.transpose(-2, -1)
            run_test(matrix, uplo)

        def run_test_skipped_elements(shape, batch, uplo):
            # check for inputs with skipped elements
            matrix = random_hermitian_matrix(shape, *batch, dtype=dtype, device=device)
            matrix = matrix[::2]
            run_test(matrix, uplo)

        shapes = (3, 5)
        batches = ((4, ), (4, 2))
        uplos = ["U", "L"]
        for shape, batch, uplo in itertools.product(shapes, batches, uplos):
            run_test_permuted(shape, batch, uplo)
            run_test_skipped_elements(shape, batch, uplo)
//  [python]
def test_eigh_non_contiguous(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_matrix

        def run_test(matrix, uplo):
            self.assertFalse(matrix.is_contiguous())
            expected_w, expected_v = np.linalg.eigh(matrix.cpu().numpy(), UPLO=uplo)
            actual_w, actual_v = torch.linalg.eigh(matrix, UPLO=uplo)
            self.assertEqual(actual_w, expected_w)
            # sign of eigenvectors is not unique and therefore absolute values are compared
            self.assertEqual(abs(actual_v), abs(expected_v))

        def run_test_permuted(shape, batch, uplo):
            # check for permuted / transposed inputs
            matrix = random_hermitian_matrix(shape, *batch, dtype=dtype, device=device)
            matrix = matrix.transpose(-2, -1)
            run_test(matrix, uplo)

        def run_test_skipped_elements(shape, batch, uplo):
            # check for inputs with skipped elements
            matrix = random_hermitian_matrix(shape, *batch, dtype=dtype, device=device)
            matrix = matrix[::2]
            run_test(matrix, uplo)

        shapes = (3, 5)
        batches = ((4, ), (4, 2))
        uplos = ["U", "L"]
        for shape, batch, uplo in itertools.product(shapes, batches, uplos):
            run_test_permuted(shape, batch, uplo)
            run_test_skipped_elements(shape, batch, uplo)
//  [python]
def test_eigh_autograd(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_matrix

        def func(x, uplo):
            x = 0.5 * (x + x.conj().transpose(-2, -1))
            return torch.symeig(x, eigenvectors=True, upper=uplo == "U")

        def func_grad_w(x, uplo):
            return func(x, uplo)[0]

        def func_grad_v(x, uplo):
            # gauge invariant loss function
            return abs(func(x, uplo)[1])

        def run_test(dims, uplo):
            x = torch.randn(*dims, dtype=dtype, device=device, requires_grad=True)

            gradcheck(func_grad_w, [x, uplo])
            gradgradcheck(func_grad_w, [x, uplo])

            gradcheck(func_grad_v, [x, uplo])
            gradgradcheck(func_grad_v, [x, uplo])

            x = random_hermitian_matrix(dims[-1], *dims[:-2]).requires_grad_()
            w, v = torch.symeig(x, eigenvectors=True)
            (w.sum() + abs(v).sum()).backward()
            self.assertEqual(x.grad, x.grad.conj().transpose(-1, -2))  # Check the gradient is Hermitian

        for dims, uplo in itertools.product([(3, 3), (2, 3, 3)], ["L", "U"]):
            run_test(dims, uplo)
//  [python]
def test_eigh_autograd(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_matrix

        def func(x, uplo):
            x = 0.5 * (x + x.conj().transpose(-2, -1))
            return torch.linalg.eigh(x, UPLO=uplo)

        def func_grad_w(x, uplo):
            return func(x, uplo)[0]

        def func_grad_v(x, uplo):
            # gauge invariant loss function
            return abs(func(x, uplo)[1])

        def run_test(dims, uplo):
            x = torch.randn(*dims, dtype=dtype, device=device, requires_grad=True)

            gradcheck(func_grad_w, [x, uplo])
            gradgradcheck(func_grad_w, [x, uplo])

            gradcheck(func_grad_v, [x, uplo])
            gradgradcheck(func_grad_v, [x, uplo])

            x = random_hermitian_matrix(dims[-1], *dims[:-2]).requires_grad_()
            w, v = torch.linalg.eigh(x)
            (w.sum() + abs(v).sum()).backward()
            self.assertEqual(x.grad, x.grad.conj().transpose(-1, -2))  # Check the gradient is Hermitian

        for dims, uplo in itertools.product([(3, 3), (2, 3, 3)], ["L", "U"]):
            run_test(dims, uplo)
//  [python]
def psd_sqrt(A: Tensor) -> Tensor:
    """Compute the unique p.s.d. square root of a positive semidefinite matrix."""
    L, U = torch.symeig(A)
    L = L[..., None, :].clamp_min(0.0)
    return U * L.sqrt() @ U.mH
//  [python]
def psd_sqrt(A: Tensor) -> Tensor:
    """Compute the unique p.s.d. square root of a positive semidefinite matrix."""
    L, U = torch.linalg.eigh(A)
    L = L[..., None, :].clamp_min(0.0)
    return U * L.sqrt() @ U.mH
//  [python]
def psd_sqrt_rsqrt(A: Tensor) -> tuple[Tensor, Tensor]:
    """Efficiently compute both the p.s.d. sqrt & pinv sqrt of p.s.d. matrix `A`."""
    L, U = torch.symeig(A)
    L = L[..., None, :].clamp_min(0.0)

    # Square root is easy
    sqrt = U * L.sqrt() @ U.mH

    # We actually compute the pseudo-inverse here for numerical stability.
    # Use the same heuristic as `torch.linalg.pinv` to determine the tolerance.
    thresh = L[..., None, -1] * A.shape[-1] * torch.finfo(A.dtype).eps
    rsqrt = U * torch.where(L > thresh, L.rsqrt(), 0.0) @ U.mH

    return sqrt, rsqrt
//  [python]
def psd_sqrt_rsqrt(A: Tensor) -> tuple[Tensor, Tensor]:
    """Efficiently compute both the p.s.d. sqrt & pinv sqrt of p.s.d. matrix `A`."""
    L, U = torch.linalg.eigh(A)
    L = L[..., None, :].clamp_min(0.0)

    # Square root is easy
    sqrt = U * L.sqrt() @ U.mH

    # We actually compute the pseudo-inverse here for numerical stability.
    # Use the same heuristic as `torch.linalg.pinv` to determine the tolerance.
    thresh = L[..., None, -1] * A.shape[-1] * torch.finfo(A.dtype).eps
    rsqrt = U * torch.where(L > thresh, L.rsqrt(), 0.0) @ U.mH

    return sqrt, rsqrt
//  [python]
def compute_eig_lapl_torch_batch(adj_matrix, pos_enc_dim=32):
    """ Compute positional encoding using graph laplacian.
        Adapted from https://github.com/graphdeeplearning/benchmarking-gnns/blob/ef8bd8c7d2c87948bc1bdd44099a52036e715cd0/data/molecules.py#L147-L168.
    
    Args:
        adj_matrix: Adjacency matrix (B x N x N).
        pos_enc_dim: Output dimensions of positional encoding.
    """
    b, n, _ = adj_matrix.size()
    
    # Laplacian
    A = adj_matrix.float()
    degree_matrix = A.sum(axis=1).clip(1)
    N = torch.diag_embed(degree_matrix ** -0.5)
    L = torch.eye(n, device=A.device)[None, ].repeat(b, 1, 1) - (N @ A) @ N

    # Eigenvectors
    eig_val, eig_vec = torch.symeig(L, eigenvectors=True)
    eig_vec = torch.flip(eig_vec, dims=[2])
    pos_enc = eig_vec[:, :, 1:pos_enc_dim + 1]

    if pos_enc.size(2) < pos_enc_dim:
        pos_enc = torch.cat([pos_enc, torch.zeros(pos_enc.size(0), pos_enc.size(1), pos_enc_dim - pos_enc.size(2), device=A.device)], dim=2)

    return pos_enc
//  [python]
def compute_eig_lapl_torch_batch(adj_matrix, pos_enc_dim=32):
    """ Compute positional encoding using graph laplacian.
        Adapted from https://github.com/graphdeeplearning/benchmarking-gnns/blob/ef8bd8c7d2c87948bc1bdd44099a52036e715cd0/data/molecules.py#L147-L168.
    
    Args:
        adj_matrix: Adjacency matrix (B x N x N).
        pos_enc_dim: Output dimensions of positional encoding.
    """
    b, n, _ = adj_matrix.size()
    
    # Laplacian
    A = adj_matrix.float()
    degree_matrix = A.sum(axis=1).clip(1)
    N = torch.diag_embed(degree_matrix ** -0.5)
    L = torch.eye(n, device=A.device)[None, ].repeat(b, 1, 1) - (N @ A) @ N

    # Eigenvectors
    eig_val, eig_vec = torch.linalg.eigh(L)
    eig_vec = torch.flip(eig_vec, dims=[2])
    pos_enc = eig_vec[:, :, 1:pos_enc_dim + 1]

    if pos_enc.size(2) < pos_enc_dim:
        pos_enc = torch.cat([pos_enc, torch.zeros(pos_enc.size(0), pos_enc.size(1), pos_enc_dim - pos_enc.size(2), device=A.device)], dim=2)

    return pos_enc
//  [python]
def optimise(self, R_init, pred_norms, pred_kappa):
        R_init = torch.from_numpy(R_init).to(self.device).to(self.dtype).t()
        pred_norms = pred_norms.to(self.dtype)

        #First normalise normal vectors to make sure they are unit vectors
        pred_norms = torch.nn.functional.normalize(pred_norms, dim=0)
        
        #If we're applying kappa downweighting then process
        if self.use_kappa:
            pred_kappa = pred_kappa.to(self.dtype)
            pred_kappa /= pred_kappa.sum(dim=1) 

        #Go through the list of required S powers and calculated their related coefficients
        #This only needs to be done once per set of normals
        S_coeffs = torch.zeros((5,5,5), device=self.device, dtype=self.dtype)

        for u,v,w in self.S_coeff_powers:
                S_coeffs[u,v,w] = self.S_func(u, v, w, pred_norms, pred_kappa)

        #Calculated final S param S_{0,0,4}, based on the other params
        S_coeffs[0,0,4] = 1 - (S_coeffs[4,0,0] + S_coeffs[0,4,0] + 2*S_coeffs[2,2,0] + 2*S_coeffs[2,0,2] + 2*S_coeffs[0,2,2])

        #Get cost function matrix from the S coefficients
        M = self.get_M_matrix(S_coeffs)
  
        #Decompose into eigenvalues (lambda, L) and other bit (Q)
        #Since M is square symmetric, M = Q @ diag(L) @ Q^t, and can use eigh instead of eig
        L, Q = torch.symeig(M)

        L = abs(L)

        # M = H^t @ H therefore H = diag(L^0.5) @ Q^t
        self.H = torch.diag(torch.sqrt(L)) @ Q.t()

        #Pass precalculated values into the LM optimisation
        R_opt, covariance = self.LM_opt(R_init)
        R_opt = torch.tensor(ScipyRot.from_matrix(R_opt).as_matrix())

        return R_opt.t(), covariance
//  [python]
def optimise(self, R_init, pred_norms, pred_kappa):
        R_init = torch.from_numpy(R_init).to(self.device).to(self.dtype).t()
        pred_norms = pred_norms.to(self.dtype)

        #First normalise normal vectors to make sure they are unit vectors
        pred_norms = torch.nn.functional.normalize(pred_norms, dim=0)
        
        #If we're applying kappa downweighting then process
        if self.use_kappa:
            pred_kappa = pred_kappa.to(self.dtype)
            pred_kappa /= pred_kappa.sum(dim=1) 

        #Go through the list of required S powers and calculated their related coefficients
        #This only needs to be done once per set of normals
        S_coeffs = torch.zeros((5,5,5), device=self.device, dtype=self.dtype)

        for u,v,w in self.S_coeff_powers:
                S_coeffs[u,v,w] = self.S_func(u, v, w, pred_norms, pred_kappa)

        #Calculated final S param S_{0,0,4}, based on the other params
        S_coeffs[0,0,4] = 1 - (S_coeffs[4,0,0] + S_coeffs[0,4,0] + 2*S_coeffs[2,2,0] + 2*S_coeffs[2,0,2] + 2*S_coeffs[0,2,2])

        #Get cost function matrix from the S coefficients
        M = self.get_M_matrix(S_coeffs)
  
        #Decompose into eigenvalues (lambda, L) and other bit (Q)
        #Since M is square symmetric, M = Q @ diag(L) @ Q^t, and can use eigh instead of eig
        L, Q = torch.linalg.eigh(M)

        L = abs(L)

        # M = H^t @ H therefore H = diag(L^0.5) @ Q^t
        self.H = torch.diag(torch.sqrt(L)) @ Q.t()

        #Pass precalculated values into the LM optimisation
        R_opt, covariance = self.LM_opt(R_init)
        R_opt = torch.tensor(ScipyRot.from_matrix(R_opt).as_matrix())

        return R_opt.t(), covariance
//  [python]
def rot_to_quat(
    rot: torch.Tensor,
) -> torch.Tensor:
    if rot.shape[-2:] != (3, 3):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot

    k = [
        [
            xx + yy + zz,
            zy - yz,
            xz - zx,
            yx - xy,
        ],
        [
            zy - yz,
            xx - yy - zz,
            xy + yx,
            xz + zx,
        ],
        [
            xz - zx,
            xy + yx,
            yy - xx - zz,
            yz + zy,
        ],
        [
            yx - xy,
            xz + zx,
            yz + zy,
            zz - xx - yy,
        ],
    ]

    k = (1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)

    _, vectors = torch.symeig(k, eigenvectors=True)
    return vectors[..., -1]
//  [python]
def rot_to_quat(
    rot: torch.Tensor,
) -> torch.Tensor:
    if rot.shape[-2:] != (3, 3):
        raise ValueError("Input rotation is incorrectly shaped")

    rot = [[rot[..., i, j] for j in range(3)] for i in range(3)]
    [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = rot

    k = [
        [
            xx + yy + zz,
            zy - yz,
            xz - zx,
            yx - xy,
        ],
        [
            zy - yz,
            xx - yy - zz,
            xy + yx,
            xz + zx,
        ],
        [
            xz - zx,
            xy + yx,
            yy - xx - zz,
            yz + zy,
        ],
        [
            yx - xy,
            xz + zx,
            yz + zy,
            zz - xx - yy,
        ],
    ]

    k = (1.0 / 3.0) * torch.stack([torch.stack(t, dim=-1) for t in k], dim=-2)

    _, vectors = torch.linalg.eigh(k)
    return vectors[..., -1]
//  [python]
def alignment(pos_y, pos_x, batch):
        with torch.no_grad():
            num_nodes = batch.n_nodes
            total_nodes = pos_y.shape[0]
            num_graphs = batch.num_graphs
            pos_y_mean = global_mean_pool(pos_y, batch.batch)
            pos_x_mean = global_mean_pool(pos_x, batch.batch)
            y = pos_y - torch.repeat_interleave(pos_y_mean, num_nodes, dim=0)
            x = pos_x - torch.repeat_interleave(pos_x_mean, num_nodes, dim=0)
            a = y + x
            b = y - x
            a = a.view(total_nodes, 1, 3)
            b = b.view(total_nodes, 3, 1)
            tmp0 = torch.cat(
                [b.new_zeros((1, 1, 1)).expand(total_nodes, -1, -1), -b.permute(0, 2, 1)], dim=-1
            )
            eye = torch.eye(3).to(a).unsqueeze(0).expand(total_nodes, -1, -1)
            a = a.expand(-1, 3, -1)
            tmp1 = torch.cross(eye, a, dim=-1)
            tmp1 = torch.cat([b, tmp1], dim=-1)
            tmp = torch.cat([tmp0, tmp1], dim=1)
            tmpb = torch.bmm(tmp.permute(0, 2, 1), tmp).view(total_nodes, -1)
            tmpb = global_mean_pool(tmpb, batch.batch).view(num_graphs, 4, 4)
            w, v = torch.symeig(tmpb, eigenvectors=True)
            min_rmsd = w[:, 0]
            min_q = v[:, :, 0]
            rotation = GNN.quaternion_to_rotation_matrix(min_q)
            t = pos_y_mean - torch.einsum("kj,kij->ki", pos_x_mean, rotation)
            rotation = torch.repeat_interleave(rotation, num_nodes, dim=0)
            t = torch.repeat_interleave(t, num_nodes, dim=0)
        pos_x = torch.einsum("kj,kij->ki", pos_x, rotation) + t
        return pos_x
//  [python]
def alignment(pos_y, pos_x, batch):
        with torch.no_grad():
            num_nodes = batch.n_nodes
            total_nodes = pos_y.shape[0]
            num_graphs = batch.num_graphs
            pos_y_mean = global_mean_pool(pos_y, batch.batch)
            pos_x_mean = global_mean_pool(pos_x, batch.batch)
            y = pos_y - torch.repeat_interleave(pos_y_mean, num_nodes, dim=0)
            x = pos_x - torch.repeat_interleave(pos_x_mean, num_nodes, dim=0)
            a = y + x
            b = y - x
            a = a.view(total_nodes, 1, 3)
            b = b.view(total_nodes, 3, 1)
            tmp0 = torch.cat(
                [b.new_zeros((1, 1, 1)).expand(total_nodes, -1, -1), -b.permute(0, 2, 1)], dim=-1
            )
            eye = torch.eye(3).to(a).unsqueeze(0).expand(total_nodes, -1, -1)
            a = a.expand(-1, 3, -1)
            tmp1 = torch.cross(eye, a, dim=-1)
            tmp1 = torch.cat([b, tmp1], dim=-1)
            tmp = torch.cat([tmp0, tmp1], dim=1)
            tmpb = torch.bmm(tmp.permute(0, 2, 1), tmp).view(total_nodes, -1)
            tmpb = global_mean_pool(tmpb, batch.batch).view(num_graphs, 4, 4)
            w, v = torch.linalg.eigh(tmpb)
            min_rmsd = w[:, 0]
            min_q = v[:, :, 0]
            rotation = GNN.quaternion_to_rotation_matrix(min_q)
            t = pos_y_mean - torch.einsum("kj,kij->ki", pos_x_mean, rotation)
            rotation = torch.repeat_interleave(rotation, num_nodes, dim=0)
            t = torch.repeat_interleave(t, num_nodes, dim=0)
        pos_x = torch.einsum("kj,kij->ki", pos_x, rotation) + t
        return pos_x
//  [python]
def auto_orient_poses(poses, method="up"):
    """Orients and centers the poses. We provide two methods for orientation: pca and up.
    pca: Orient the poses so that the principal component of the points is aligned with the axes.
        This method works well when all of the cameras are in the same plane.
    up: Orient the poses so that the average up vector is aligned with the z axis.
        This method works well when images are not at arbitrary angles.
    Args:
        poses: The poses to orient.
        method: The method to use for orientation. Either "pca" or "up".
    Returns:
        The oriented poses.
    borrowed from from nerfstudio
    """
    translation = poses[..., :3, 3]

    mean_translation = torch.mean(translation, dim=0)
    translation = translation - mean_translation

    if method == "pca":
        _, eigvec = torch.symeig(translation.T @ translation, eigenvectors=True)
        eigvec = torch.flip(eigvec, dims=(-1,))

        if torch.det(eigvec) < 0:
            eigvec[:, 2] = -eigvec[:, 2]

        transform = torch.cat([eigvec, eigvec @ -mean_translation[..., None]], dim=-1)
        oriented_poses = transform @ poses

        if oriented_poses.mean(axis=0)[2, 1] < 0:
            oriented_poses[:, 1:3] = -1 * oriented_poses[:, 1:3]
    elif method == "up":
        up = torch.mean(poses[:, :3, 1], dim=0)
        up = up / torch.norm(up)

        rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))
        transform = torch.cat([rotation, rotation @ -mean_translation[..., None]], dim=-1)
        oriented_poses = transform @ poses

    return oriented_poses
//  [python]
def auto_orient_poses(poses, method="up"):
    """Orients and centers the poses. We provide two methods for orientation: pca and up.
    pca: Orient the poses so that the principal component of the points is aligned with the axes.
        This method works well when all of the cameras are in the same plane.
    up: Orient the poses so that the average up vector is aligned with the z axis.
        This method works well when images are not at arbitrary angles.
    Args:
        poses: The poses to orient.
        method: The method to use for orientation. Either "pca" or "up".
    Returns:
        The oriented poses.
    borrowed from from nerfstudio
    """
    translation = poses[..., :3, 3]

    mean_translation = torch.mean(translation, dim=0)
    translation = translation - mean_translation

    if method == "pca":
        _, eigvec = torch.linalg.eigh(translation.T @ translation)
        eigvec = torch.flip(eigvec, dims=(-1,))

        if torch.linalg.det(eigvec) < 0:
            eigvec[:, 2] = -eigvec[:, 2]

        transform = torch.cat([eigvec, eigvec @ -mean_translation[..., None]], dim=-1)
        oriented_poses = transform @ poses

        if oriented_poses.mean(axis=0)[2, 1] < 0:
            oriented_poses[:, 1:3] = -1 * oriented_poses[:, 1:3]
    elif method == "up":
        up = torch.mean(poses[:, :3, 1], dim=0)
        up = up / torch.linalg.norm(up)

        rotation = rotation_matrix(up, torch.Tensor([0, 0, 1]))
        transform = torch.cat([rotation, rotation @ -mean_translation[..., None]], dim=-1)
        oriented_poses = transform @ poses

    return oriented_poses
//  [python]
def alignment_loss(pos_y, pos_x, batch, clamp=None):
        with torch.no_grad():
            num_nodes = batch.n_nodes
            total_nodes = pos_y.shape[0]
            num_graphs = batch.num_graphs
            pos_y_mean = global_mean_pool(pos_y, batch.batch)
            pos_x_mean = global_mean_pool(pos_x, batch.batch)
            y = pos_y - torch.repeat_interleave(pos_y_mean, num_nodes, dim=0)
            x = pos_x - torch.repeat_interleave(pos_x_mean, num_nodes, dim=0)
            a = y + x
            b = y - x
            a = a.view(total_nodes, 1, 3)
            b = b.view(total_nodes, 3, 1)
            tmp0 = torch.cat(
                [b.new_zeros((1, 1, 1)).expand(total_nodes, -1, -1), -b.permute(0, 2, 1)], dim=-1
            )
            eye = torch.eye(3).to(a).unsqueeze(0).expand(total_nodes, -1, -1)
            a = a.expand(-1, 3, -1)
            tmp1 = torch.cross(eye, a, dim=-1)
            tmp1 = torch.cat([b, tmp1], dim=-1)
            tmp = torch.cat([tmp0, tmp1], dim=1)
            tmpb = torch.bmm(tmp.permute(0, 2, 1), tmp).view(total_nodes, -1)
            tmpb = global_mean_pool(tmpb, batch.batch).view(num_graphs, 4, 4)
            w, v = torch.symeig(tmpb, eigenvectors=True)
            min_rmsd = w[:, 0]
            min_q = v[:, :, 0]
            rotation = GNN.quaternion_to_rotation_matrix(min_q)
            t = pos_y_mean - torch.einsum("kj,kij->ki", pos_x_mean, rotation)
            rotation = torch.repeat_interleave(rotation, num_nodes, dim=0)
            t = torch.repeat_interleave(t, num_nodes, dim=0)
        pos_x = torch.einsum("kj,kij->ki", pos_x, rotation) + t
        if clamp is None:
            loss = global_mean_pool((pos_y - pos_x).norm(dim=-1, keepdim=True), batch.batch).mean()
        else:
            loss = global_mean_pool((pos_y - pos_x).norm(dim=-1, keepdim=True), batch.batch)
            loss = torch.clamp(loss, min=clamp).mean()
        return loss, min_rmsd.mean()
//  [python]
def alignment_loss(pos_y, pos_x, batch, clamp=None):
        with torch.no_grad():
            num_nodes = batch.n_nodes
            total_nodes = pos_y.shape[0]
            num_graphs = batch.num_graphs
            pos_y_mean = global_mean_pool(pos_y, batch.batch)
            pos_x_mean = global_mean_pool(pos_x, batch.batch)
            y = pos_y - torch.repeat_interleave(pos_y_mean, num_nodes, dim=0)
            x = pos_x - torch.repeat_interleave(pos_x_mean, num_nodes, dim=0)
            a = y + x
            b = y - x
            a = a.view(total_nodes, 1, 3)
            b = b.view(total_nodes, 3, 1)
            tmp0 = torch.cat(
                [b.new_zeros((1, 1, 1)).expand(total_nodes, -1, -1), -b.permute(0, 2, 1)], dim=-1
            )
            eye = torch.eye(3).to(a).unsqueeze(0).expand(total_nodes, -1, -1)
            a = a.expand(-1, 3, -1)
            tmp1 = torch.cross(eye, a, dim=-1)
            tmp1 = torch.cat([b, tmp1], dim=-1)
            tmp = torch.cat([tmp0, tmp1], dim=1)
            tmpb = torch.bmm(tmp.permute(0, 2, 1), tmp).view(total_nodes, -1)
            tmpb = global_mean_pool(tmpb, batch.batch).view(num_graphs, 4, 4)
            w, v = torch.linalg.eigh(tmpb)
            min_rmsd = w[:, 0]
            min_q = v[:, :, 0]
            rotation = GNN.quaternion_to_rotation_matrix(min_q)
            t = pos_y_mean - torch.einsum("kj,kij->ki", pos_x_mean, rotation)
            rotation = torch.repeat_interleave(rotation, num_nodes, dim=0)
            t = torch.repeat_interleave(t, num_nodes, dim=0)
        pos_x = torch.einsum("kj,kij->ki", pos_x, rotation) + t
        if clamp is None:
            loss = global_mean_pool((pos_y - pos_x).norm(dim=-1, keepdim=True), batch.batch).mean()
        else:
            loss = global_mean_pool((pos_y - pos_x).norm(dim=-1, keepdim=True), batch.batch)
            loss = torch.clamp(loss, min=clamp).mean()
        return loss, min_rmsd.mean()
//  [python]
def pseudo_points(self):
        if not isinstance(self._variational_distribution, CholeskyVariationalDistribution):
            raise NotImplementedError('Only CholeskyVariationalDistribution has pseudo-point support currently, ', 'but your _variational_distribution is a ', self._variational_distribution.__name__)
        var_cov_root = TriangularLinearOperator(self._variational_distribution.chol_variational_covar)
        var_cov = CholLinearOperator(var_cov_root)
        var_mean = self.variational_distribution.mean
        if var_mean.shape[-1] != 1:
            var_mean = var_mean.unsqueeze(-1)
        Kmm = self.model.covar_module(self.inducing_points)
        res = Kmm - var_cov
        cov_diff = res
        eval_lhs = var_cov.to_dense()
        eval_rhs = cov_diff.transpose(-1, -2).matmul(eval_lhs)
        inner_term = cov_diff.matmul(cov_diff.transpose(-1, -2))
        inner_solve = inner_term.add_jitter(self.jitter_val).solve(eval_rhs, eval_lhs.transpose(-1, -2))
        inducing_covar = var_cov + inner_solve
        rhs = cov_diff.transpose(-1, -2).matmul(var_mean)
        inner_rhs_mean_solve = inner_term.add_jitter(self.jitter_val).solve(rhs)
        pseudo_target_mean = var_mean + var_cov.matmul(inner_rhs_mean_solve)
        try:
            pseudo_target_covar = CholLinearOperator(inducing_covar.add_jitter(self.jitter_val).cholesky()).to_dense()
        except NotPSDError:
            evals, evecs = torch.symeig(inducing_covar)
            pseudo_target_covar = evecs.matmul(DiagLinearOperator(evals + self.jitter_val)).matmul(evecs.transpose(-1, -2)).to_dense()
        return pseudo_target_covar, pseudo_target_mean
//  [python]
def pseudo_points(self):
        if not isinstance(self._variational_distribution, CholeskyVariationalDistribution):
            raise NotImplementedError('Only CholeskyVariationalDistribution has pseudo-point support currently, ', 'but your _variational_distribution is a ', self._variational_distribution.__name__)
        var_cov_root = TriangularLinearOperator(self._variational_distribution.chol_variational_covar)
        var_cov = CholLinearOperator(var_cov_root)
        var_mean = self.variational_distribution.mean
        if var_mean.shape[-1] != 1:
            var_mean = var_mean.unsqueeze(-1)
        Kmm = self.model.covar_module(self.inducing_points)
        res = Kmm - var_cov
        cov_diff = res
        eval_lhs = var_cov.to_dense()
        eval_rhs = cov_diff.transpose(-1, -2).matmul(eval_lhs)
        inner_term = cov_diff.matmul(cov_diff.transpose(-1, -2))
        inner_solve = inner_term.add_jitter(self.jitter_val).solve(eval_rhs, eval_lhs.transpose(-1, -2))
        inducing_covar = var_cov + inner_solve
        rhs = cov_diff.transpose(-1, -2).matmul(var_mean)
        inner_rhs_mean_solve = inner_term.add_jitter(self.jitter_val).solve(rhs)
        pseudo_target_mean = var_mean + var_cov.matmul(inner_rhs_mean_solve)
        try:
            pseudo_target_covar = CholLinearOperator(inducing_covar.add_jitter(self.jitter_val).cholesky()).to_dense()
        except NotPSDError:
            evals, evecs = torch.linalg.eigh(inducing_covar)
            pseudo_target_covar = evecs.matmul(DiagLinearOperator(evals + self.jitter_val)).matmul(evecs.transpose(-1, -2)).to_dense()
        return pseudo_target_covar, pseudo_target_mean
//  [python]
def post_covariance_training(self, **kwargs):
        # Calculate top right singular vectors from covariance matrices
        self.top_singular_vectors = {
            k: torch.symeig(cov, eigenvectors=True).eigenvectors[:, -1]
            for k, cov in self.covariances.items()
        }
//  [python]
def post_covariance_training(self, **kwargs):
        # Calculate top right singular vectors from covariance matrices
        self.top_singular_vectors = {
            k: torch.linalg.eigh(cov).eigenvectors[:, -1]
            for k, cov in self.covariances.items()
        }
//  [python]
def symeig(A: Tensor, largest: Optional[bool] = False) -> Tuple[Tensor, Tensor]:
    """Return eigenpairs of A with specified ordering.
    """
    if largest is None:
        largest = False
    E, Z = torch.symeig(A, eigenvectors=True, upper=True)
    # assuming that E is ordered
    if largest:
        E = torch.flip(E, dims=(-1,))
        Z = torch.flip(Z, dims=(-1,))
    return E, Z
//  [python]
def symeig(A: Tensor, largest: Optional[bool] = False) -> Tuple[Tensor, Tensor]:
    """Return eigenpairs of A with specified ordering.
    """
    if largest is None:
        largest = False
    E, Z = torch.linalg.eigh(A, UPLO='U')
    # assuming that E is ordered
    if largest:
        E = torch.flip(E, dims=(-1,))
        Z = torch.flip(Z, dims=(-1,))
    return E, Z
//  [python]
def decompose_hessian(hessian, lengths, shapes):
    H_matrix = torch.zeros(sum(lengths), sum(lengths))
    for i, j in product(range(len(shapes)), range(len(shapes))):
        i0, i1 = sum(lengths[:i+1]) - lengths[i], sum(lengths[:i+1])
        j0, j1 = sum(lengths[:j+1]) - lengths[j], sum(lengths[:j+1])
        H_matrix[i0:i1, j0:j1] = hessian[i][j].reshape(lengths[i], lengths[j])
    vals, vecs = torch.symeig(H_matrix)
    vs = []
    for i in range(sum(lengths)):
        vec = vecs[:, i]
        v = tuple(vec[sum(lengths[:i]):sum(lengths[:i+1])].reshape(shapes[i]) for i in range(len(lengths)))
        vs.append(v)
    return vals, vs
//  [python]
def decompose_hessian(hessian, lengths, shapes):
    H_matrix = torch.zeros(sum(lengths), sum(lengths))
    for i, j in product(range(len(shapes)), range(len(shapes))):
        i0, i1 = sum(lengths[:i+1]) - lengths[i], sum(lengths[:i+1])
        j0, j1 = sum(lengths[:j+1]) - lengths[j], sum(lengths[:j+1])
        H_matrix[i0:i1, j0:j1] = hessian[i][j].reshape(lengths[i], lengths[j])
    vals, vecs = torch.linalg.eigh(H_matrix)
    vs = []
    for i in range(sum(lengths)):
        vec = vecs[:, i]
        v = tuple(vec[sum(lengths[:i]):sum(lengths[:i+1])].reshape(shapes[i]) for i in range(len(lengths)))
        vs.append(v)
    return vals, vs