jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def apply_affine2point(points, theta, shape):
  assert points.size(0) == 3, 'invalid points shape : {:}'.format(points.size())
  with torch.no_grad():
    ok_points = points[2,:] == 1
    assert torch.sum(ok_points).item() > 0, 'there is no visiable point'
    points[:2,:] = normalize_points(shape, points[:2,:])

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

    trans_points = torch.linalg.solve(theta, points[:, ok_points])
    norm_trans_points[:, ok_points] = trans_points
    
  return norm_trans_points
//  [python]
def test_solve(self, device, dtype):
        for (k, n) in zip([2, 3, 5], [3, 5, 7]):
            b, A = self.solve_test_helper((n,), (n, k), device, dtype)
            x = torch.solve(b, A)[0]
            self.assertLessEqual(b.dist(A.mm(x)), 1e-12)
//  [python]
def test_solve(self, device, dtype):
        for (k, n) in zip([2, 3, 5], [3, 5, 7]):
            b, A = self.solve_test_helper((n,), (n, k), device, dtype)
            x = torch.linalg.solve(A, b)
            self.assertLessEqual(b.dist(A.mm(x)), 1e-12)
//  [python]
def find_coeffs(p1, p2):
    "Find coefficients for warp tfm from `p1` to `p2`"
    m = []
    p = p1[:,0,0]
    #The equations we'll need to solve.
    for i in range(p1.shape[1]):
        m.append(stack([p2[:,i,0], p2[:,i,1], t1(p), t0(p), t0(p), t0(p), -p1[:,i,0]*p2[:,i,0], -p1[:,i,0]*p2[:,i,1]]))
        m.append(stack([t0(p), t0(p), t0(p), p2[:,i,0], p2[:,i,1], t1(p), -p1[:,i,1]*p2[:,i,0], -p1[:,i,1]*p2[:,i,1]]))
    #The 8 scalars we seek are solution of AX = B
    A = stack(m).permute(2, 0, 1)
    B = p1.view(p1.shape[0], 8, 1)
    return torch.solve(B,A)[0]
//  [python]
def find_coeffs(p1, p2):
    "Find coefficients for warp tfm from `p1` to `p2`"
    m = []
    p = p1[:,0,0]
    #The equations we'll need to solve.
    for i in range(p1.shape[1]):
        m.append(stack([p2[:,i,0], p2[:,i,1], t1(p), t0(p), t0(p), t0(p), -p1[:,i,0]*p2[:,i,0], -p1[:,i,0]*p2[:,i,1]]))
        m.append(stack([t0(p), t0(p), t0(p), p2[:,i,0], p2[:,i,1], t1(p), -p1[:,i,1]*p2[:,i,0], -p1[:,i,1]*p2[:,i,1]]))
    #The 8 scalars we seek are solution of AX = B
    A = stack(m).permute(2, 0, 1)
    B = p1.view(p1.shape[0], 8, 1)
    return torch.linalg.solve(A, B)
//  [python]
def _solve_P_Q(U, V, structure=None):
    """
    A helper function for expm_2009.
    Parameters
    ----------
    U : ndarray
        Pade numerator.
    V : ndarray
        Pade denominator.
    structure : str, optional
        A string describing the structure of both matrices `U` and `V`.
        Only `upper_triangular` is currently supported.
    Notes
    -----
    The `structure` argument is inspired by similar args
    for theano and cvxopt functions.
    """
    P = U + V
    Q = -U + V
    return torch.solve(P, Q)[0]
//  [python]
def _solve_P_Q(U, V, structure=None):
    """
    A helper function for expm_2009.
    Parameters
    ----------
    U : ndarray
        Pade numerator.
    V : ndarray
        Pade denominator.
    structure : str, optional
        A string describing the structure of both matrices `U` and `V`.
        Only `upper_triangular` is currently supported.
    Notes
    -----
    The `structure` argument is inspired by similar args
    for theano and cvxopt functions.
    """
    P = U + V
    Q = -U + V
    return torch.linalg.solve(P, Q)
//  [python]
def test_solve_batched(self, device, dtype):
        def solve_batch_helper(A_dims, b_dims):
            b, A = self.solve_test_helper(A_dims, b_dims, device, dtype)
            x_exp_list = []
            for i in range(b_dims[0]):
                x_exp_list.append(torch.solve(b[i], A[i])[0])
            x_exp = torch.stack(x_exp_list)  # Stacked output
            x_act = torch.solve(b, A)[0]  # Actual output
            self.assertEqual(x_exp, x_act)  # Equality check
            self.assertLessEqual(b.dist(torch.matmul(A, x_act)), 1e-12)  # Correctness check

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

        for batchsize in [1, 3, 4]:
            solve_batch_helper((5, batchsize), (batchsize, 5, 10))
//  [python]
def test_solve_batched_many_batches(self, device, dtype):
        b, A = self.solve_test_helper((5, 256, 256), (5, 1), device, dtype)
        x, _ = torch.solve(b, A)
        self.assertEqual(torch.matmul(A, x), b.expand(A.shape[:-2] + (5, 1)))

        b, A = self.solve_test_helper((3,), (512, 512, 3, 1), device, dtype)
        x, _ = torch.solve(b, A)
        self.assertEqual(torch.matmul(A, x), b)
//  [python]
def test_solve_batched_many_batches(self, device, dtype):
        b, A = self.solve_test_helper((5, 256, 256), (5, 1), device, dtype)
        x = torch.linalg.solve(A, b)
        self.assertEqual(torch.matmul(A, x), b.expand(A.shape[:-2] + (5, 1)))

        b, A = self.solve_test_helper((3,), (512, 512, 3, 1), device, dtype)
        x = torch.linalg.solve(A, b)
        self.assertEqual(torch.matmul(A, x), b)
//  [python]
def test_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value
        A = random_fullrank_matrix_distinct_singular_value(2, 2, dtype=dtype,
                                                           device=device).permute(1, 0, 2)
        b = torch.randn(2, 2, 2, dtype=dtype, device=device).permute(2, 1, 0)
        x, _ = torch.solve(b, A)
        x_exp = torch.Tensor(solve(A.cpu().numpy(), b.cpu().numpy())).to(dtype=dtype, device=device)
        self.assertEqual(x, x_exp)
//  [python]
def test_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_fullrank_matrix_distinct_singular_value
        A = random_fullrank_matrix_distinct_singular_value(2, 2, dtype=dtype,
                                                           device=device).permute(1, 0, 2)
        b = torch.randn(2, 2, 2, dtype=dtype, device=device).permute(2, 1, 0)
        x, _ = torch.linalg.solve(A, b)
        x_exp = torch.Tensor(solve(A.cpu().numpy(), b.cpu().numpy())).to(dtype=dtype, device=device)
        self.assertEqual(x, x_exp)
//  [python]
def test_solve_batched_broadcasting(self, device, dtype):
        from numpy.linalg import solve

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

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

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

        # test against numpy.linalg.solve
        run_test((2, 1, 3, 4, 4), (2, 1, 3, 4, 6))  # no broadcasting
        run_test((2, 1, 3, 4, 4), (4, 6))  # broadcasting b
        run_test((4, 4), (2, 1, 3, 4, 2))  # broadcasting A
        run_test((1, 3, 1, 4, 4), (2, 1, 3, 4, 5))
//  [python]
def fit(self, x, y, z):
        x, y, z = (torch.Tensor(v.flatten()) for v in [x, y, z])
        A, AT = self._construct_A(x, y, z**2)
        coeffs = torch.solve(AT @ z[...,None], AT @ A)[0]
        self.b  = coeffs[0][0]
        self.ai = coeffs[1:].flatten()
//  [python]
def fit(self, x, y, z):
        x, y, z = (torch.Tensor(v.flatten()) for v in [x, y, z])
        A, AT = self._construct_A(x, y, z**2)
        coeffs = torch.linalg.solve(AT @ A, AT @ z[...,None])
        self.b  = coeffs[0][0]
        self.ai = coeffs[1:].flatten()
//  [python]
def b_inv(b_mat):
    """code from https://stackoverflow.com/questions/46595157/how-to-apply-the-
    torch-inverse-function-of-pytorch-to-every-sample-in-the-batc.

    :param b_mat:
    :return:
    """
    eye = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    try:
        b_inv, _ = torch.solve(eye, b_mat)
    except:
        # https://github.com/zju3dv/clean-pvnet/issues/8
        b_inv = eye
    return b_inv
//  [python]
def b_inv(b_mat):
    eye = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    try:
        b_inv = torch.linalg.solve(eye, b_mat)
    except:
        b_inv = eye
    return b_inv
//  [python]
def loglikelihood(self, value, verbose=False):
        err = (value - self.mean).view(-1,1).contiguous()
        covar = gpytorch.add_jitter( self.covariance_matrix , jitter_val=1*1e-8)
        quatterm = err.T @ torch.solve(err, covar).solution
        logdetcovar = torch.logdet(covar)
        ll = - 0.5*quatterm - 0.5*logdetcovar - 0.5*len(covar)*math.log(2*math.pi)
        if verbose:
            print(f'log|K|: {logdetcovar.item():10.2f} {" ":5s} |err|^2_{{K^-1}}: {quatterm.item():10.2f} {" ":5s}',end='')
        return ll
//  [python]
def loglikelihood(self, value, verbose=False):
        err = (value - self.mean).view(-1,1).contiguous()
        covar = gpytorch.add_jitter( self.covariance_matrix , jitter_val=1*1e-8)
        quatterm = err.T @ torch.linalg.solve(covar, err)
        logdetcovar = torch.logdet(covar)
        ll = - 0.5*quatterm - 0.5*logdetcovar - 0.5*len(covar)*math.log(2*math.pi)
        if verbose:
            print(f'log|K|: {logdetcovar.item():10.2f} {" ":5s} |err|^2_{{K^-1}}: {quatterm.item():10.2f} {" ":5s}',end='')
        return ll
//  [python]
def test_solve_removed_error(self, device):
        a = make_tensor(5, 5, device=device, dtype=torch.float32)
        b = make_tensor(5, 1, device=device, dtype=torch.float32)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            torch.solve(b, a)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            b.solve(a)
//  [python]
def test_solve_removed_error(self, device):
        a = make_tensor(5, 5, device=device, dtype=torch.float32)
        b = make_tensor(5, 1, device=device, dtype=torch.float32)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            torch.linalg.solve(a, b)
        with self.assertRaisesRegex(RuntimeError, "This function was deprecated since version 1.9 and is now removed"):
            b.solve(a)
//  [python]
def torch_expm(A):
    n_A = A.shape[0]
    A_fro = torch.sqrt(A.abs().pow(2).sum(dim=(1, 2), keepdim=True))
    maxnorm = torch.tensor([5.371920351148152], dtype=A.dtype, device=A.device)
    zero = torch.tensor([0.0], dtype=A.dtype, device=A.device)
    n_squarings = torch.max(zero, torch.ceil(torch_log2(A_fro / maxnorm)))
    A_scaled = A / 2.0 ** n_squarings
    n_squarings = n_squarings.flatten().type(torch.int64)
    U, V = torch_pade13(A_scaled)
    P = U + V
    Q = -U + V
    R, _ = torch.solve(P, Q)
    res = [R]
    for i in range(int(n_squarings.max())):
        res.append(res[-1].matmul(res[-1]))
    R = torch.stack(res)
    expmA = R[n_squarings, torch.arange(n_A)]
    return expmA[0]
//  [python]
def torch_expm(A):
    n_A = A.shape[0]
    A_fro = torch.sqrt(A.abs().pow(2).sum(dim=(1, 2), keepdim=True))
    maxnorm = torch.tensor([5.371920351148152], dtype=A.dtype, device=A.device)
    zero = torch.tensor([0.0], dtype=A.dtype, device=A.device)
    n_squarings = torch.max(zero, torch.ceil(torch_log2(A_fro / maxnorm)))
    A_scaled = A / 2.0 ** n_squarings
    n_squarings = n_squarings.flatten().type(torch.int64)
    U, V = torch_pade13(A_scaled)
    P = U + V
    Q = -U + V
    R, _ = torch.linalg.solve(P, Q)
    res = [R]
    for i in range(int(n_squarings.max())):
        res.append(res[-1].matmul(res[-1]))
    R = torch.stack(res)
    expmA = R[n_squarings, torch.arange(n_A)]
    return expmA[0]
//  [python]
def test_solve_methods_arg_device(self, device):
        for b_device, A_device in product(['cpu', device], repeat=2):
            if b_device == A_device:
                continue

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

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

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

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

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

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

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

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

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

            # This checks if a suitable error message is thrown
            # when LU output and pivots are on the same device
            with self.assertRaisesRegex(RuntimeError,
                                        "Expected LU_pivots and LU_data to be on the same device"):
                torch.lu_solve(b, A, torch.rand(A.shape[:-1], device=b_device).int())
//  [python]
def solve_linear_system(A, b, sparse=False):
    return torch.solve(b, A).solution
//  [python]
def solve_linear_system(A, b, sparse=False):
    return torch.linalg.solve(A, b)
//  [python]
def fast_global_registration(xyz, xyz_target, mu_init=1, num_iter=20):
    """
    estimate the rotation and translation using Fast Global Registration algorithm (M estimator for robust estimation)
    http://vladlen.info/papers/fast-global-registration.pdf
    """
    assert xyz.shape == xyz_target.shape

    T_res = torch.eye(4, device=xyz.device)
    mu = mu_init
    source = xyz.clone()
    weight = torch.ones(len(source), 1, device=xyz.device)
    for i in range(num_iter):
        if i > 0 and i % 5 == 0:
            mu /= 2.0
        A, b = get_matrix_system(source, xyz_target, weight)
        x, _ = torch.solve(A.T @ b, A.T.mm(A))
        T = get_trans(x.view(-1))
        source = source.mm(T[:3, :3].T) + T[:3, 3]
        T_res = T @ T_res
        weight = get_geman_mclure_weight(source, xyz_target, mu)
    return T_res
//  [python]
def fast_global_registration(xyz, xyz_target, mu_init=1, num_iter=20):
    """
    estimate the rotation and translation using Fast Global Registration algorithm (M estimator for robust estimation)
    http://vladlen.info/papers/fast-global-registration.pdf
    """
    assert xyz.shape == xyz_target.shape

    T_res = torch.eye(4, device=xyz.device)
    mu = mu_init
    source = xyz.clone()
    weight = torch.ones(len(source), 1, device=xyz.device)
    for i in range(num_iter):
        if i > 0 and i % 5 == 0:
            mu /= 2.0
        A, b = get_matrix_system(source, xyz_target, weight)
        x = torch.linalg.solve(A.T @ A, A.T @ b)
        T = get_trans(x.view(-1))
        source = source.mm(T[:3, :3].T) + T[:3, 3]
        T_res = T @ T_res
        weight = get_geman_mclure_weight(source, xyz_target, mu)
    return T_res
//  [python]
def b_inv(b_mat):
    '''
    code from
    https://stackoverflow.com/questions/46595157/how-to-apply-the-torch-inverse-function-of-pytorch-to-every-sample-in-the-batc
    :param b_mat:
    :return:
    '''
    eye = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    if torch.__version__ >= '1.0.0':
        b_inv, _ = torch.solve(eye, b_mat)
    else:
        b_inv, _ = torch.gesv(eye, b_mat)
    # b_inv, _ = torch.gesv(eye, b_mat)
    return b_inv
//  [python]
def b_inv(b_mat):
    '''
    code from
    https://stackoverflow.com/questions/46595157/how-to-apply-the-torch-inverse-function-of-pytorch-to-every-sample-in-the-batc
    :param b_mat:
    :return:
    '''
    eye = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    b_inv = torch.linalg.solve(b_mat, eye)
    return b_inv
//  [python]
def batch_lstsq(
    input: torch.Tensor,  # matrix B of shape (batch * m * k) 
    A: torch.Tensor  # matrix A of shape (batch * m * n) 
):

    # X = torch.bmm(
    #     torch.pinverse(A),
    #     input
    # )

    AT = A.permute(0,2,1)
    
    X = torch.solve(
        torch.bmm(AT, input),
        torch.bmm(AT, A)
    ).solution

    # X = torch.bmm(
    #     torch.inverse(torch.bmm(AT, A)),
    #     torch.bmm(AT, input)
    # )
    
    return X
//  [python]
def batch_lstsq(
    input: torch.Tensor,  # matrix B of shape (batch * m * k) 
    A: torch.Tensor  # matrix A of shape (batch * m * n) 
):

    AT = A.permute(0,2,1)
    
    X = torch.linalg.solve(
        torch.bmm(AT, A),
        torch.bmm(AT, input)
    )
    
    return X
//  [python]
def find_coeffs(orig_pts, targ_pts):
    matrix = []
    #The equations we'll need to solve.
    for p1, p2 in zip(targ_pts, orig_pts):
        matrix.append([p1[0], p1[1], 1, 0, 0, 0, -p2[0]*p1[0], -p2[0]*p1[1]])
        matrix.append([0, 0, 0, p1[0], p1[1], 1, -p2[1]*p1[0], -p2[1]*p1[1]])

    A = FloatTensor(matrix)
    B = FloatTensor(orig_pts).view(8, 1)
    #The 8 scalars we seek are solution of AX = B
    return list(torch.solve(B,A)[0][:,0])
//  [python]
def find_coeffs(orig_pts, targ_pts):
    matrix = []
    #The equations we'll need to solve.
    for p1, p2 in zip(targ_pts, orig_pts):
        matrix.append([p1[0], p1[1], 1, 0, 0, 0, -p2[0]*p1[0], -p2[0]*p1[1]])
        matrix.append([0, 0, 0, p1[0], p1[1], 1, -p2[1]*p1[0], -p2[1]*p1[1]])

    A = FloatTensor(matrix)
    B = FloatTensor(orig_pts).view(8, 1)
    #The 8 scalars we seek are solution of AX = B
    return list(torch.linalg.solve(A, B)[:,0])
//  [python]
def BMRW(x, y, w):
    x_norm = calculate_l1_norm(x)
    y_norm = calculate_l1_norm(y)
    eye_x = torch.eye(x.size(1)).float().to(x.device)

    latent_z = F.softmax(torch.einsum('nkd,ntd->nkt', [y_norm, x_norm]) * 5.0, 1)
    norm_latent_z = latent_z / (latent_z.sum(dim=-1, keepdim=True) + 1e-9)
    affinity_mat = torch.einsum('nkt,nkd->ntd', [latent_z, norm_latent_z])
    # mat_inv_x, _ = torch.linalg.solve(eye_x, eye_x - (w ** 2) * affinity_mat)
    mat_inv_x, _ = torch.solve(eye_x, eye_x - (w ** 2) * affinity_mat)
    y2x_sum_x = w * torch.einsum('nkt,nkd->ntd', [latent_z, y]) + x
    refined_x = (1 - w) * torch.einsum('ntk,nkd->ntd', [mat_inv_x, y2x_sum_x])    

    return refined_x
//  [python]
def BMRW(x, y, w):
    x_norm = calculate_l1_norm(x)
    y_norm = calculate_l1_norm(y)
    eye_x = torch.eye(x.size(1)).float().to(x.device)

    latent_z = F.softmax(torch.einsum('nkd,ntd->nkt', [y_norm, x_norm]) * 5.0, 1)
    norm_latent_z = latent_z / (latent_z.sum(dim=-1, keepdim=True) + 1e-9)
    affinity_mat = torch.einsum('nkt,nkd->ntd', [latent_z, norm_latent_z])
    mat_inv_x, _ = torch.linalg.solve(eye_x, eye_x - (w ** 2) * affinity_mat)
    # mat_inv_x, _ = torch.solve(eye_x, eye_x - (w ** 2) * affinity_mat)
    y2x_sum_x = w * torch.einsum('nkt,nkd->ntd', [latent_z, y]) + x
    refined_x = (1 - w) * torch.einsum('ntk,nkd->ntd', [mat_inv_x, y2x_sum_x])    

    return refined_x
//  [python]
def log(self, x, y):
        ytx = tb.xty(y, x)
        At = tb.transpose(y) - ytx @ tb.transpose(x)
        Bt = torch.solve(At, ytx).solution
        us, ss, vs = tb.svd(tb.transpose(Bt))
        return tb.mvmt(us[..., :self.p],
                       ss[..., :self.p].atan(),
                       vs[..., :self.p])
//  [python]
def log(self, x, y):
        ytx = tb.xty(y, x)
        At = tb.transpose(y) - ytx @ tb.transpose(x)
        Bt = torch.linalg.solve(At, ytx)
        us, ss, vs = tb.svd(tb.transpose(Bt))
        return tb.mvmt(us[..., :self.p],
                       ss[..., :self.p].atan(),
                       vs[..., :self.p])
//  [python]
def _stable_solve(b, a, eps=1e-6):
    try:
        return torch.solve(b, a)[0]
    except RuntimeError:
        a = condition_scm(a, eps)
        return torch.solve(b, a)[0]
//  [python]
def _stable_solve(b, a, eps=1e-6):
    try:
        return torch.linalg.solve(a, b)
    except RuntimeError:
        a = condition_scm(a, eps)
        return torch.linalg.solve(a, b)
//  [python]
def solve(matrix, rhs):
    return torch.solve(matrix, rhs)
//  [python]
def solve(matrix, rhs):
    return torch.linalg.solve(matrix, rhs)
//  [python]
def M(self, q, eps=1e-6, verbose=False):
        """Returns a function that multiplies the mass matrix M(q) by a vector qdot
        Args:
            q: bs x D Tensor representing the position
        """
        assert q.ndim == 2
        lower_triangular = self.tril_Minv(q) * 0
        assert lower_triangular.ndim == 3

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

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

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

        return M_func
//  [python]
def update(self, i, u, cov, zupt):
        H = self.H
        if zupt == 1:
            z = 1
        else:
            z = 0

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

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

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

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

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

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

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

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

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

        R = torch.diag(torch.cat((cov,
            self.zupt_omega_cov,
            self.zupt_acc_cov), 0))
        S = axat(H, self.P) + R
        Kt, _ = torch.linalg.solve(S, self.P.mm(H.t()).t())
        K = Kt.t()
        self.xi = K.mv(self.r)
        self.state_update(i)
        self.covariance_update(i, K, H, R)
//  [python]
def compute_forward_dynamics(
        self,
        q: torch.Tensor,
        qd: torch.Tensor,
        f: torch.Tensor,
        include_gravity: Optional[bool] = True,
    ) -> torch.Tensor:
        r"""
        Computes next qdd by solving the Euler-Lagrange equation
        qdd = H^{-1} (F - Cv - G - damping_term)

        Args:
            q: joint angles [batch_size x n_dofs]
            qd: joint velocities [batch_size x n_dofs]
            f: forces to be applied [batch_size x n_dofs]
            include_gravity: set to False if your robot has gravity compensation

        Returns: accelerations that are the result of applying forces f in state q, qd

        """

        nle = self.compute_non_linear_effects(
            q=q, qd=qd, include_gravity=include_gravity
        )
        inertia_mat = self.compute_lagrangian_inertia_matrix(
            q=q, include_gravity=include_gravity
        )

        # Solve H qdd = F - Cv - G - damping_term
        qdd = torch.solve(f.unsqueeze(2) - nle.unsqueeze(2), inertia_mat)[0].squeeze(2)

        return qdd
//  [python]
def compute_forward_dynamics(
        self,
        q: torch.Tensor,
        qd: torch.Tensor,
        f: torch.Tensor,
        include_gravity: Optional[bool] = True,
    ) -> torch.Tensor:
        r"""
        Computes next qdd by solving the Euler-Lagrange equation
        qdd = H^{-1} (F - Cv - G - damping_term)

        Args:
            q: joint angles [batch_size x n_dofs]
            qd: joint velocities [batch_size x n_dofs]
            f: forces to be applied [batch_size x n_dofs]
            include_gravity: set to False if your robot has gravity compensation

        Returns: accelerations that are the result of applying forces f in state q, qd

        """

        nle = self.compute_non_linear_effects(
            q=q, qd=qd, include_gravity=include_gravity
        )
        inertia_mat = self.compute_lagrangian_inertia_matrix(
            q=q, include_gravity=include_gravity
        )

        # Solve H qdd = F - Cv - G - damping_term
        qdd = torch.linalg.solve(inertia_mat, f.unsqueeze(2) - nle.unsqueeze(2)).squeeze(2)

        return qdd
//  [python]
def UCB(self, A, phi):
        tmp, LU = torch.solve(phi, A)
        return torch.sqrt(torch.mm(torch.transpose(phi, 0, 1), tmp))
//  [python]
def UCB(self, A, phi):
        tmp, LU = torch.linalg.solve(A, phi)
        return torch.sqrt(torch.mm(torch.transpose(phi, 0, 1), tmp))
//  [python]
def fit(self):
        if self.readout_training in {'gd', 'svd'}:
            return
        if self.readout_training == 'cholesky':
            W = torch.solve(self.XTy, self.XTX + self.lambda_reg * torch.eye(self.XTX.size(0), device=self.XTX.device))[0].t()
            self.XTX = None
            self.XTy = None
            self.readout.bias = nn.Parameter(W[:, 0])
            self.readout.weight = nn.Parameter(W[:, 1:])
        elif self.readout_training == 'inv':
            I = self.lambda_reg * torch.eye(self.XTX.size(0))
            A = self.XTX + I
            X_rank = torch.linalg.matrix_rank(A).item()
            if X_rank == self.X.size(0):
                W = torch.mm(torch.inverse(A), self.XTy).t()
            else:
                W = torch.mm(torch.pinverse(A), self.XTy).t()
            self.readout.bias = nn.Parameter(W[:, 0])
            self.readout.weight = nn.Parameter(W[:, 1:])
            self.XTX = None
            self.XTy = None
//  [python]
def fit(self):
        if self.readout_training in {'gd', 'svd'}:
            return
        if self.readout_training == 'cholesky':
            W = torch.linalg.solve(self.XTX + self.lambda_reg * torch.eye(self.XTX.size(0), device=self.XTX.device), self.XTy).t()
            self.XTX = None
            self.XTy = None
            self.readout.bias = nn.Parameter(W[:, 0])
            self.readout.weight = nn.Parameter(W[:, 1:])
        elif self.readout_training == 'inv':
            I = self.lambda_reg * torch.eye(self.XTX.size(0))
            A = self.XTX + I
            X_rank = torch.linalg.matrix_rank(A).item()
            if X_rank == self.X.size(0):
                W = torch.mm(torch.inverse(A), self.XTy).t()
            else:
                W = torch.mm(torch.pinverse(A), self.XTy).t()
            self.readout.bias = nn.Parameter(W[:, 0])
            self.readout.weight = nn.Parameter(W[:, 1:])
            self.XTX = None
            self.XTy = None
//  [python]
def receive_reward(self, arm, context, reward):
        self.bb += reward * self.last_feature
        theta, LU = torch.solve(self.bb, self.Design)
        with torch.no_grad():
            self.model.fc.weight.copy_(theta.view(1, -1))
        self.collector.collect_data(context, arm, reward)
        self.last_cxt = context
//  [python]
def receive_reward(self, arm, context, reward):
        self.bb += reward * self.last_feature
        theta, LU = torch.linalg.solve(self.Design, self.bb)
        with torch.no_grad():
            self.model.fc.weight.copy_(theta.view(1, -1))
        self.collector.collect_data(context, arm, reward)
        self.last_cxt = context
//  [python]
def forward(self, x, index=None):
        # similarity between sleep and wake data
        G = self.kernel(self.x_fit, x)
        if self.training:
            self.update_K()

        # this computes G.t() @ K^{-1}
        GKinv = torch.solve(G, self.K)[0].t()

        if not self.centre_elbo:
            if index is None:
                return [GKinv @ f for f in self.fs]
            else:
                return GKinv @ self.fs[index]
        else:
            if index is None:
                elbo_mean = self.fs[-1].mean(dim=0)
                return [GKinv @ f for f in self.fs[:-1]] + [GKinv @ (self.fs[-1] - elbo_mean) + elbo_mean]
            else:
                if index == -1 or index == len(self.fs) - 1:
                    elbo_mean = self.fs[-1].mean(dim=0)
                    return GKinv @ (self.fs[index] - elbo_mean) + elbo_mean
                else:
                    return GKinv @ self.fs[index]
//  [python]
def forward(self, x, index=None):
        # similarity between sleep and wake data
        G = self.kernel(self.x_fit, x)
        if self.training:
            self.update_K()

        # this computes G.t() @ K^{-1}
        GKinv = torch.linalg.solve(self.K, G).t()

        if not self.centre_elbo:
            if index is None:
                return [GKinv @ f for f in self.fs]
            else:
                return GKinv @ self.fs[index]
        else:
            if index is None:
                elbo_mean = self.fs[-1].mean(dim=0)
                return [GKinv @ f for f in self.fs[:-1]] + [GKinv @ (self.fs[-1] - elbo_mean) + elbo_mean]
            else:
                if index == -1 or index == len(self.fs) - 1:
                    elbo_mean = self.fs[-1].mean(dim=0)
                    return GKinv @ (self.fs[index] - elbo_mean) + elbo_mean
                else:
                    return GKinv @ self.fs[index]
//  [python]
def fit_full_poly(self, epsilon=1, k=2, lmbd=0, target='hypo', fkine=None):
        X = self.support_points
        if fkine is not None:
            X = fkine(X).reshape([len(X), -1])
            self.fkine = fkine
            self.support_fkine = X
        self.poly_kernel = kernel.Polyharmonic(k=k, epsilon=epsilon)
        phi = self.poly_kernel(X, X)
        phi.fill_diagonal_(lmbd)
        print(phi.shape)
        l1 = torch.cat([phi, X, torch.ones((len(X), 1))], dim=1)
        l2 = torch.cat([X.T, torch.zeros((X.shape[1], X.shape[1]+1))], dim=1)
        l3 = torch.cat([torch.ones((1, len(X))), torch.zeros(1, X.shape[1]+1)], dim=1)
        print([l.shape for l in [l1, l2, l3]])
        L = torch.cat([l1, l2, l3], dim=0)
        if target == 'hypo':
            y = self.hypothesis
        elif 'dist' in target:
            y = self.distance
        elif 'label' in target:
            y = self.y
        self.poly_nodes = torch.solve(
            torch.cat([y, torch.zeros(X.shape[1]+1, self.num_class)], dim=0),
            L).solution
//  [python]
def fit_full_poly(self, epsilon=1, k=2, lmbd=0, target='hypo', fkine=None):
        X = self.support_points
        if fkine is not None:
            X = fkine(X).reshape([len(X), -1])
            self.fkine = fkine
            self.support_fkine = X
        self.poly_kernel = kernel.Polyharmonic(k=k, epsilon=epsilon)
        phi = self.poly_kernel(X, X)
        phi.fill_diagonal_(lmbd)
        print(phi.shape)
        l1 = torch.cat([phi, X, torch.ones((len(X), 1))], dim=1)
        l2 = torch.cat([X.T, torch.zeros((X.shape[1], X.shape[1]+1))], dim=1)
        l3 = torch.cat([torch.ones((1, len(X))), torch.zeros(1, X.shape[1]+1)], dim=1)
        print([l.shape for l in [l1, l2, l3]])
        L = torch.cat([l1, l2, l3], dim=0)
        if target == 'hypo':
            y = self.hypothesis
        elif 'dist' in target:
            y = self.distance
        elif 'label' in target:
            y = self.y
        self.poly_nodes = torch.linalg.solve(
            torch.cat([y, torch.zeros(X.shape[1]+1, self.num_class)], dim=0),
            L)
//  [python]
def stabilization(DPhi,Phi):
    DPhiT = DPhi.transpose(-1, -2)
    X,_ = torch.solve(Phi.unsqueeze(-1),DPhiT @ J(DPhi))
    return -J(DPhi@X).squeeze(-1)
//  [python]
def stabilization(DPhi,Phi):
    DPhiT = DPhi.transpose(-1, -2)
    X,_ = torch.linalg.solve(DPhiT @ J(DPhi),Phi.unsqueeze(-1))
    return -J(DPhi@X).squeeze(-1)
//  [python]
def one_adapt_step(
        self,
        model: torch.nn.Module,
        optimizer: torch.optim.Optimizer,
        batch: Batch,
        timer: Timer,
        random_seed: int = None,
    ):
        """adapt the model in one step."""
        eps = self._meta_conf.model_eps
        num_classes = self._meta_conf.statistics["n_classes"]
        with timer("forward"):
            with fork_rng_with_seed(random_seed):
                y_hat = model(batch._x)
            y_hat = y_hat / self._meta_conf.temperature_scaling
            softmax_prob = F.softmax(y_hat, dim=1)

            # adapt.
            smax_inp = softmax_prob
            eye = torch.eye(num_classes).to(y_hat.device)
            eye = eye.reshape((1, num_classes, num_classes))
            eye = eye.repeat(y_hat.shape[0], 1, 1)
            t2 = eps * torch.diag_embed(smax_inp)
            smax_inp = torch.unsqueeze(smax_inp, 2)
            t3 = eps * torch.bmm(smax_inp, torch.transpose(smax_inp, 1, 2))
            matrix = eye + t2 - t3
            # y_star = torch.linalg.solve(matrix, smax_inp)  # only support in a latest torch version.
            y_star, _ = torch.solve(smax_inp, matrix)
            y_star = torch.squeeze(y_star)

            pseudo_prob = y_star
            loss = torch.logsumexp(y_hat, dim=1) - (
                pseudo_prob * y_hat - eps * pseudo_prob * (1 - softmax_prob)
            ).sum(dim=1)
            loss = loss.mean()

            # apply fisher regularization when enabled
            if self.fishers is not None:
                ewc_loss = 0
                for name, param in model.named_parameters():
                    if name in self.fishers:
                        ewc_loss += (
                            self._meta_conf.fisher_alpha
                            * (
                                self.fishers[name][0]
                                * (param - self.fishers[name][1]) ** 2
                            ).sum()
                        )
                loss += ewc_loss

        with timer("backward"):
            loss.backward()
            grads = dict(
                (name, param.grad.clone().detach())
                for name, param in model.named_parameters()
                if param.grad is not None
            )
            optimizer.step()
            optimizer.zero_grad()
        return {
            "optimizer": copy.deepcopy(optimizer).state_dict(),
            "loss": loss.item(),
            "grads": grads,
            "yhat": y_hat,
        }
//  [python]
def one_adapt_step(
        self,
        model: torch.nn.Module,
        optimizer: torch.optim.Optimizer,
        batch: Batch,
        timer: Timer,
        random_seed: int = None,
    ):
        """adapt the model in one step."""
        eps = self._meta_conf.model_eps
        num_classes = self._meta_conf.statistics["n_classes"]
        with timer("forward"):
            with fork_rng_with_seed(random_seed):
                y_hat = model(batch._x)
            y_hat = y_hat / self._meta_conf.temperature_scaling
            softmax_prob = F.softmax(y_hat, dim=1)

            # adapt.
            smax_inp = softmax_prob
            eye = torch.eye(num_classes).to(y_hat.device)
            eye = eye.reshape((1, num_classes, num_classes))
            eye = eye.repeat(y_hat.shape[0], 1, 1)
            t2 = eps * torch.diag_embed(smax_inp)
            smax_inp = torch.unsqueeze(smax_inp, 2)
            t3 = eps * torch.bmm(smax_inp, torch.transpose(smax_inp, 1, 2))
            matrix = eye + t2 - t3
            y_star = torch.linalg.solve(matrix, smax_inp)
            y_star = torch.squeeze(y_star)

            pseudo_prob = y_star
            loss = torch.logsumexp(y_hat, dim=1) - (
                pseudo_prob * y_hat - eps * pseudo_prob * (1 - softmax_prob)
            ).sum(dim=1)
            loss = loss.mean()

            # apply fisher regularization when enabled
            if self.fishers is not None:
                ewc_loss = 0
                for name, param in model.named_parameters():
                    if name in self.fishers:
                        ewc_loss += (
                            self._meta_conf.fisher_alpha
                            * (
                                self.fishers[name][0]
                                * (param - self.fishers[name][1]) ** 2
                            ).sum()
                        )
                loss += ewc_loss

        with timer("backward"):
            loss.backward()
            grads = dict(
                (name, param.grad.clone().detach())
                for name, param in model.named_parameters()
                if param.grad is not None
            )
            optimizer.step()
            optimizer.zero_grad()
        return {
            "optimizer": copy.deepcopy(optimizer).state_dict(),
            "loss": loss.item(),
            "grads": grads,
            "yhat": y_hat,
        }
//  [python]
def Proj(DPhi):
    if DPhi.shape[-1]==0: return lambda M:M # (no constraints)
    def _P(M):
        DPhiT = DPhi.transpose(-1, -2)
        X, _ = torch.solve(DPhiT @ M, DPhiT @ J(DPhi))
        return M - J(DPhi @ X)

    return _P
//  [python]
def Proj(DPhi):
    if DPhi.shape[-1]==0: return lambda M:M # (no constraints)
    def _P(M):
        DPhiT = DPhi.transpose(-1, -2)
        X = torch.linalg.solve(DPhiT @ J(DPhi), DPhiT @ M)
        return M - J(DPhi @ X)

    return _P
//  [python]
def basis_transform(
        self,
        new_basis: Tensor,
        ERROR_EPS_ADD: float = 0.0,
        ERROR_EPS_MUL: float = 0.0,
    ) -> Tuple["HybridZonotope", Tensor]:
        if self.errors is None:  # No errors to transform in a different basis
            return self, torch.ones_like(self.head) - 1e-8
        # We solve for the coordinates (x) of curr_basis (B) in new_basis (A)
        # I.e. we solve Ax=b
        A = new_basis
        B = torch.flatten(self.errors, start_dim=1).T

        device = self.device
        dtype = self.dtype

        if A.shape[0] < 500 or A.shape[0] != A.shape[1]:
            # depending on the size of the matrices different methods are faster
            if isinstance(A, torch.Tensor):
                A = A.cpu().detach().numpy()
            B = B.cpu().detach().numpy()
            sol = np.linalg.lstsq(A, B, rcond=None)[0]
            sol = torch.tensor(sol, dtype=dtype, device=device)
        else:
            if not isinstance(A, torch.Tensor):
                A = torch.tensor(A)
            sol = torch.solve(B, A).solution

        assert sol is not None, "No solution found"
        assert torch.isclose(
            torch.matmul(A, sol), B, atol=1e-7, rtol=1e-6
        ).all(), "Projection into new base errors failed"

        # We add the component ERROR_EPS_ADD to ensure the resulting error matrix has full rank and to compensate for potential numerical errors
        x = torch.sum(sol.abs(), dim=1) * (1 + ERROR_EPS_MUL) + ERROR_EPS_ADD

        new_errors = (
            torch.tensor(x.reshape(1, -1) * new_basis, dtype=dtype)
            .T.unsqueeze(1)
            .view(-1, *self.head.shape)
        )

        return HybridZonotope(self.head, self.beta, new_errors, self.domain), x
//  [python]
def basis_transform(
        self,
        new_basis: Tensor,
        ERROR_EPS_ADD: float = 0.0,
        ERROR_EPS_MUL: float = 0.0,
    ) -> Tuple["HybridZonotope", Tensor]:
        if self.errors is None:  # No errors to transform in a different basis
            return self, torch.ones_like(self.head) - 1e-8
        # We solve for the coordinates (x) of curr_basis (B) in new_basis (A)
        # I.e. we solve Ax=b
        A = new_basis
        B = torch.flatten(self.errors, start_dim=1).T

        device = self.device
        dtype = self.dtype

        if A.shape[0] < 500 or A.shape[0] != A.shape[1]:
            # depending on the size of the matrices different methods are faster
            if isinstance(A, torch.Tensor):
                A = A.cpu().detach().numpy()
            B = B.cpu().detach().numpy()
            sol = np.linalg.lstsq(A, B, rcond=None)[0]
            sol = torch.tensor(sol, dtype=dtype, device=device)
        else:
            if not isinstance(A, torch.Tensor):
                A = torch.tensor(A)
            sol = torch.linalg.solve(A, B)

        assert sol is not None, "No solution found"
        assert torch.isclose(
            torch.matmul(A, sol), B, atol=1e-7, rtol=1e-6
        ).all(), "Projection into new base errors failed"

        # We add the component ERROR_EPS_ADD to ensure the resulting error matrix has full rank and to compensate for potential numerical errors
        x = torch.sum(sol.abs(), dim=1) * (1 + ERROR_EPS_MUL) + ERROR_EPS_ADD

        new_errors = (
            torch.tensor(x.reshape(1, -1) * new_basis, dtype=dtype)
            .T.unsqueeze(1)
            .view(-1, *self.head.shape)
        )

        return HybridZonotope(self.head, self.beta, new_errors, self.domain), x
//  [python]
def test_distance_formulas(seed, rand_spd, d):
    spd = SPD(d)
    x, y = rand_spd(2, d)
    ref_dist = spd.dist(x, y)

    # compute :math:`Y^{-1} X` and take its eigenvalues (we have to use
    # `torch.eig` for this as the resulting matrix might not be symmetric)
    d1 = torch.solve(y, x)[0].eig()[0][:, 0].log_().pow_(2).sum().sqrt_()
    assert_allclose(ref_dist, d1, atol=1e-4)

    d2 = torch.solve(x, y)[0].eig()[0][:, 0].log_().pow_(2).sum().sqrt_()
    assert_allclose(ref_dist, d2, atol=1e-4)
//  [python]
def test_distance_formulas(seed, rand_spd, d):
    spd = SPD(d)
    x, y = rand_spd(2, d)
    ref_dist = spd.dist(x, y)

    # compute :math:`Y^{-1} X` and take its eigenvalues (we have to use
    # `torch.eig` for this as the resulting matrix might not be symmetric)
    d1 = torch.linalg.solve(y, x).eig()[0][:, 0].log_().pow_(2).sum().sqrt_()
    assert_allclose(ref_dist, d1, atol=1e-4)

    d2 = torch.linalg.solve(x, y).eig()[0][:, 0].log_().pow_(2).sum().sqrt_()
    assert_allclose(ref_dist, d2, atol=1e-4)
//  [python]
def binv(b_mat):
    """
    Computes an inverse of each matrix in the batch.
    Pytorch 0.4.1 does not support batched matrix inverse.
    Hence, we are solving AX=I.

    Parameters:
      b_mat:  a (n_batch, n, n) Tensor.
    Returns: a (n_batch, n, n) Tensor.
    """
    id_matrix = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    b_inv, _ = torch.solve(id_matrix, b_mat)
    return b_inv
//  [python]
def binv(b_mat):
    """
    Computes an inverse of each matrix in the batch.
    Pytorch 0.4.1 does not support batched matrix inverse.
    Hence, we are solving AX=I.

    Parameters:
      b_mat:  a (n_batch, n, n) Tensor.
    Returns: a (n_batch, n, n) Tensor.
    """
    id_matrix = b_mat.new_ones(b_mat.size(-1)).diag().expand_as(b_mat)
    b_inv = torch.linalg.solve(b_mat, id_matrix)
    return b_inv
//  [python]
def solve(b: ComplexTensor, a: ComplexTensor, return_LU=False) -> ComplexTensor:
    """Solve ax = b"""
    a = complex_matrix2real_matrix(a)
    b = complex_vector2real_vector(b)
    if LooseVersion(torch.__version__) >= LooseVersion("1.8"):
        if return_LU:
            LU, pivots = torch.lu(a)
            x = torch.lu_solve(b, LU, pivots)
        else:
            x = torch.linalg.solve(a, b)
    else:
        x, LU = torch.solve(b, a)
    if return_LU:
        return real_vector2complex_vector(x), real_matrix2complex_matrix(LU)
    else:
        return real_vector2complex_vector(x)
//  [python]
def solve(b: ComplexTensor, a: ComplexTensor, return_LU=False) -> ComplexTensor:
    """Solve ax = b"""
    a = complex_matrix2real_matrix(a)
    b = complex_vector2real_vector(b)
    if LooseVersion(torch.__version__) >= LooseVersion("1.8"):
        if return_LU:
            LU, pivots = torch.lu(a)
            x = torch.lu_solve(b, LU, pivots)
        else:
            x = torch.linalg.solve(a, b)
    else:
        x = torch.linalg.solve(a, b)
    if return_LU:
        return real_vector2complex_vector(x), real_matrix2complex_matrix(LU)
    else:
        return real_vector2complex_vector(x)
//  [python]
def test_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from common_utils import random_fullrank_matrix_distinct_singular_value
        A = random_fullrank_matrix_distinct_singular_value(2, 2, dtype=dtype,
                                                           device=device).permute(1, 0, 2)
        b = torch.randn(2, 2, 2, dtype=dtype, device=device).permute(2, 1, 0)
        x, _ = torch.solve(b, A)
        x_exp = torch.Tensor(solve(A.cpu().numpy(), b.cpu().numpy())).to(dtype=dtype, device=device)
        self.assertEqual(x.data, x_exp)