jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_cholesky(self):
        linear_operator = self.create_linear_operator()
        evaluated = self.evaluate_linear_operator(linear_operator)
        for upper in (False, True):
            res = torch.cholesky(linear_operator, upper=upper)
            actual = torch.cholesky(evaluated, upper=upper)
            self.assertAllClose(res.to_dense(), actual, rtol=1e-3, atol=1e-5)

            rhs = torch.randn(linear_operator.size(-2), 3)
            self.assertAllClose(
                torch.cholesky_solve(rhs, res, upper=upper),
                rhs.cholesky_solve(actual, upper=upper),
                rtol=1e-3,
                atol=1e-5,
            )
//  [python]
def test_cholesky(self):
        linear_operator = self.create_linear_operator()
        evaluated = self.evaluate_linear_operator(linear_operator)
        for upper in (False, True):
            res = torch.linalg.cholesky(linear_operator, upper=upper)
            actual = torch.linalg.cholesky(evaluated, upper=upper)
            self.assertAllClose(res.to_dense(), actual, rtol=1e-3, atol=1e-5)

            rhs = torch.randn(linear_operator.size(-2), 3)
            self.assertAllClose(
                torch.linalg.cholesky_solve(rhs, res, upper=upper),
                rhs.cholesky_solve(actual, upper=upper),
                rtol=1e-3,
                atol=1e-5,
            )
//  [python]
def mc_sample_fast(mu, cov, n_sample):
    r_sample = torch.randn((n_sample,) + mu.shape, dtype=mu.dtype, device=mu.device)
    sample = mu + (torch.cholesky(cov) @ r_sample.unsqueeze(dim=-1)).squeeze(dim=-1)
    return sample
//  [python]
def mc_sample_fast(mu, cov, n_sample):
    r_sample = torch.randn((n_sample,) + mu.shape, dtype=mu.dtype, device=mu.device)
    sample = mu + (torch.linalg.cholesky(cov) @ r_sample.unsqueeze(dim=-1)).squeeze(dim=-1)
    return sample
//  [python]
def test_batch_cg_with_tridiag(self):
        batch = 5
        size = 10
        matrix = torch.randn(batch, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.transpose(-1, -2))
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        rhs = torch.randn(batch, size, 10, dtype=torch.float64)
        solves, t_mats = linear_cg(matrix.matmul,
                                   rhs=rhs,
                                   n_tridiag=8,
                                   max_iter=size,
                                   max_tridiag_iter=10,
                                   tolerance=0,
                                   eps=1e-30)

        # Check cg
        matrix_chol = torch.cholesky(matrix)
        actual = torch.cholesky_solve(rhs, matrix_chol)
        assert (torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))

        # Check tridiag
        for i in range(5):
            eigs = matrix[i].symeig()[0]
            for j in range(8):
                approx_eigs = t_mats[j, i].symeig()[0]
                assert (torch.allclose(eigs, approx_eigs, atol=1e-3, rtol=1e-4))
//  [python]
def test_batch_cg_with_tridiag(self):
        batch = 5
        size = 10
        matrix = torch.randn(batch, size, size, dtype=torch.float64)
        matrix = matrix.matmul(matrix.transpose(-1, -2))
        matrix.div_(matrix.norm())
        matrix.add_(torch.eye(matrix.size(-1), dtype=torch.float64).mul_(1e-1))

        rhs = torch.randn(batch, size, 10, dtype=torch.float64)
        solves, t_mats = linear_cg(matrix.matmul,
                                   rhs=rhs,
                                   n_tridiag=8,
                                   max_iter=size,
                                   max_tridiag_iter=10,
                                   tolerance=0,
                                   eps=1e-30)

        # Check cg
        matrix_chol = torch.linalg.cholesky(matrix)
        actual = torch.linalg.solve_triangular(rhs, matrix_chol, upper=False)
        assert (torch.allclose(solves, actual, atol=1e-3, rtol=1e-4))

        # Check tridiag
        for i in range(5):
            eigs = matrix[i].symeig()[0]
            for j in range(8):
                approx_eigs = t_mats[j, i].symeig()[0]
                assert (torch.allclose(eigs, approx_eigs, atol=1e-3, rtol=1e-4))
//  [python]
def cholesky(x):
    return torch.cholesky(x)
//  [python]
def cholesky(x):
    return torch.linalg.cholesky(x)
//  [python]
def psd_safe_cholesky(A, upper=False, out=None, jitter=None):
    """Compute the Cholesky decomposition of A. If A is only p.s.d, add a small jitter to the diagonal.
    Args:
        :attr:`A` (Tensor):
            The tensor to compute the Cholesky decomposition of
        :attr:`upper` (bool, optional):
            See torch.cholesky
        :attr:`out` (Tensor, optional):
            See torch.cholesky
        :attr:`jitter` (float, optional):
            The jitter to add to the diagonal of A in case A is only p.s.d. If omitted, chosen
            as 1e-6 (float) or 1e-8 (double)
    """
    try:
        if cg.constant_jitter is not None:
            A.diagonal(dim1=-2, dim2=-1).add_(cg.constant_jitter)

        L = torch.cholesky(A, upper=upper, out=out)

        ## For a weird reason sometimes even A has nan torch.cholesky doesnt fail and goes into except. We check it manually and raise it ourselves
        isnan = torch.isnan(A)
        if isnan.any():
            raise NanError(
                f"cholesky_cpu: {isnan.sum().item()} of {A.numel()} elements of the {A.shape} tensor are NaN."
            )

        return L, A
    except RuntimeError as e:
        isnan = torch.isnan(A)
        if isnan.any():
            raise NanError(
                f"cholesky_cpu: {isnan.sum().item()} of {A.numel()} elements of the {A.shape} tensor are NaN."
            )

        if jitter is None:
            jitter = 1e-6 if A.dtype == torch.float32 else 1e-8
        Aprime = A.clone()
        jitter_prev = 0
        for i in range(3):
            jitter_new = jitter * (10 ** i)
            Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
            jitter_prev = jitter_new
            try:
                L = torch.cholesky(Aprime, upper=upper, out=out)
                warnings.warn(f"A not p.d., added jitter of {jitter_new} to the diagonal", NumericalWarning)
                return L, Aprime
            except RuntimeError:
                continue
        raise e
//  [python]
def psd_safe_cholesky(A, upper=False, out=None, jitter=None):
    """Compute the Cholesky decomposition of A. If A is only p.s.d, add a small jitter to the diagonal.
    Args:
        :attr:`A` (Tensor):
            The tensor to compute the Cholesky decomposition of
        :attr:`upper` (bool, optional):
            See torch.linalg.cholesky
        :attr:`out` (Tensor, optional):
            See torch.linalg.cholesky
        :attr:`jitter` (float, optional):
            The jitter to add to the diagonal of A in case A is only p.s.d. If omitted, chosen
            as 1e-6 (float) or 1e-8 (double)
    """
    try:
        if cg.constant_jitter is not None:
            A.diagonal(dim1=-2, dim2=-1).add_(cg.constant_jitter)

        L = torch.linalg.cholesky(A, upper=upper, out=out)

        ## For a weird reason sometimes even A has nan torch.cholesky doesnt fail and goes into except. We check it manually and raise it ourselves
        isnan = torch.isnan(A)
        if isnan.any():
            raise NanError(
                f"cholesky_cpu: {isnan.sum().item()} of {A.numel()} elements of the {A.shape} tensor are NaN."
            )

        return L, A
    except RuntimeError as e:
        isnan = torch.isnan(A)
        if isnan.any():
            raise NanError(
                f"cholesky_cpu: {isnan.sum().item()} of {A.numel()} elements of the {A.shape} tensor are NaN."
            )

        if jitter is None:
            jitter = 1e-6 if A.dtype == torch.float32 else 1e-8
        Aprime = A.clone()
        jitter_prev = 0
        for i in range(3):
            jitter_new = jitter * (10 ** i)
            Aprime.diagonal(dim1=-2, dim2=-1).add_(jitter_new - jitter_prev)
            jitter_prev = jitter_new
            try:
                L = torch.linalg.cholesky(Aprime, upper=upper, out=out)
                warnings.warn(f"A not p.d., added jitter of {jitter_new} to the diagonal", NumericalWarning)
                return L, Aprime
            except RuntimeError:
                continue
        raise e
//  [python]
def qmc_sample_fast(mu, cov, n_sample, rng):
    qr_seq = torch.stack([box_muller_transform(rng.draw(n_sample)) for _ in range(mu.size(0))], dim=1).unsqueeze(dim=2).type_as(mu)
    sample = mu + (torch.cholesky(cov) @ qr_seq.unsqueeze(dim=-1)).squeeze(dim=-1)
    return sample
//  [python]
def qmc_sample_fast(mu, cov, n_sample, rng):
    qr_seq = torch.stack([box_muller_transform(rng.draw(n_sample)) for _ in range(mu.size(0))], dim=1).unsqueeze(dim=2).type_as(mu)
    sample = mu + (torch.linalg.cholesky(cov) @ qr_seq.unsqueeze(dim=-1)).squeeze(dim=-1)
    return sample
//  [python]
def purposive_sample_fast(mu, cov, n_sample, loc_sample):
    loc_norm = box_muller_transform(loc_sample).permute(2, 0, 1, 3).expand((n_sample,) + mu.shape)
    sample = mu + (torch.cholesky(cov) @ loc_norm.unsqueeze(dim=-1)).squeeze(dim=-1)
    return sample
//  [python]
def purposive_sample_fast(mu, cov, n_sample, loc_sample):
    loc_norm = box_muller_transform(loc_sample).permute(2, 0, 1, 3).expand((n_sample,) + mu.shape)
    sample = mu + (torch.linalg.cholesky(cov) @ loc_norm.unsqueeze(dim=-1)).squeeze(dim=-1)
    return sample
//  [python]
def log_trace(x):
    x = torch.cholesky(x)
    diag = torch.diagonal(x, dim1=-2, dim2=-1)
    return 2 * torch.sum(torch.log(diag + 1e-8), dim=-1)
//  [python]
def log_trace(x):
    x = torch.linalg.cholesky(x)
    diag = torch.diagonal(x, dim1=-2, dim2=-1)
    return 2 * torch.sum(torch.log(diag + 1e-8), dim=-1)
//  [python]
def __init__(self, units=64):
        """Initializes a dynamics model for our door task."""

        super().__init__(state_dim=3)

        control_dim = 7

        # Fixed dynamics covariance
        self.Q_scale_tril = nn.Parameter(
            torch.cholesky(torch.diag(torch.FloatTensor([0.05, 0.01, 0.01]))),
            requires_grad=False,
        )

        # Build neural network
        self.state_layers = layers.state_layers(units=units)
        self.control_layers = layers.control_layers(units=units)
        self.shared_layers = nn.Sequential(
            nn.Linear(units * 2, units),
            resblocks.Linear(units),
            resblocks.Linear(units),
            resblocks.Linear(units),
            nn.Linear(units, self.state_dim + 1),
        )
        self.units = units
//  [python]
def __init__(self, units=64):
        """Initializes a dynamics model for our door task."""

        super().__init__(state_dim=3)

        control_dim = 7

        # Fixed dynamics covariance
        self.Q_scale_tril = nn.Parameter(
            torch.linalg.cholesky(torch.diag(torch.FloatTensor([0.05, 0.01, 0.01]))),
            requires_grad=False,
        )

        # Build neural network
        self.state_layers = layers.state_layers(units=units)
        self.control_layers = layers.control_layers(units=units)
        self.shared_layers = nn.Sequential(
            nn.Linear(units * 2, units),
            resblocks.Linear(units),
            resblocks.Linear(units),
            resblocks.Linear(units),
            nn.Linear(units, self.state_dim + 1),
        )
        self.units = units
//  [python]
def _socp_safety(self, cbc2, x, u0,
                     safety_factor=None,
                     extravars=None):
        """
        Var(CBC2) = Au² + b' u + c
        E(CBC2) = e' u + e
        """
        factor = safety_factor
        m = u0.shape[-1]

        (bfe, e), (V, bfv, v), mean, var = cbc2_quadratic_terms(cbc2, x, u0)
        with torch.no_grad():
            # [1, u] Asq [1; u]
            Asq = torch.cat(
                (
                    torch.cat((torch.tensor([[v]]),         (bfv / 2).reshape(1, -1)), dim=-1),
                    torch.cat(((bfv / 2).reshape(-1, 1),    V), dim=-1)
                ),
                dim=-2)

            # [1, u] Asq [1; u] = |L[1; u]|_2 = |[0, A] [y_1; y_2; u] + b|_2
            A = torch.zeros((m + 1, m + extravars))
            try:
                L = torch.cholesky(Asq) # (m+1) x (m+1)
            except RuntimeError as err:
                if "cholesky" in str(err) and "singular" in str(err):
                    diag_e, V = torch.symeig(Asq, eigenvectors=True)
                    L = torch.max(torch.diag(diag_e),
                                    torch.tensor(0.)).sqrt() @ V.t()
                else:
                    raise
            A[:, extravars:] = L[:, 1:]
            b = L[:, 0] # (m+1)
            c = torch.zeros((m+extravars,))
            c[extravars:] = bfe
            # # We want to return in format?
            # (name, (A, b, c, d))
            # s.t. factor * ||A[y_1; u] + b||_2 <= c'x + d
            return (factor * A, factor * b, c, e)
//  [python]
def _socp_safety(self, cbc2, x, u0,
                     safety_factor=None,
                     extravars=None):
        """
        Var(CBC2) = Au² + b' u + c
        E(CBC2) = e' u + e
        """
        factor = safety_factor
        m = u0.shape[-1]

        (bfe, e), (V, bfv, v), mean, var = cbc2_quadratic_terms(cbc2, x, u0)
        with torch.no_grad():
            # [1, u] Asq [1; u]
            Asq = torch.cat(
                (
                    torch.cat((torch.tensor([[v]]),         (bfv / 2).reshape(1, -1)), dim=-1),
                    torch.cat(((bfv / 2).reshape(-1, 1),    V), dim=-1)
                ),
                dim=-2)

            # [1, u] Asq [1; u] = |L[1; u]|_2 = |[0, A] [y_1; y_2; u] + b|_2
            A = torch.zeros((m + 1, m + extravars))
            try:
                L = torch.linalg.cholesky(Asq) # (m+1) x (m+1)
            except RuntimeError as err:
                if "cholesky" in str(err) and "singular" in str(err):
                    diag_e, V = torch.symeig(Asq, eigenvectors=True)
                    L = torch.max(torch.diag(diag_e),
                                    torch.tensor(0.)).sqrt() @ V.t()
                else:
                    raise
            A[:, extravars:] = L[:, 1:]
            b = L[:, 0] # (m+1)
            c = torch.zeros((m+extravars,))
            c[extravars:] = bfe
            # # We want to return in format?
            # (name, (A, b, c, d))
            # s.t. factor * ||A[y_1; u] + b||_2 <= c'x + d
            return (factor * A, factor * b, c, e)
//  [python]
def forward(self, x_train, y_train, x_test):
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        K_s = self.kernel(x_train, x_test)
        K_ss = self.kernel(x_test, x_test)
        
        # recommended implementation, fastest so far
        L = torch.cholesky(K)
        Alpha = torch.cholesky_solve(y_train, L)
        mu = K_s.T @ Alpha
        # v = torch.cholesky_solve(K_s, L)    # wrong implementation
        v = L.inverse() @ K_s   # correct implementation
        cov = K_ss - v.T @ v

        cov = cov.diag().view(-1, 1).expand_as(mu)
        return mu.squeeze(), cov
//  [python]
def forward(self, x_train, y_train, x_test):
        K = self.kernel(x_train, x_train) + self.noise_variance.pow(2) * torch.eye(len(x_train))
        K_s = self.kernel(x_train, x_test)
        K_ss = self.kernel(x_test, x_test)
        
        # recommended implementation, fastest so far
        L = torch.linalg.cholesky(K)
        Alpha = torch.linalg.solve_triangular(L, y_train)
        mu = K_s.T @ Alpha
        # v = torch.cholesky_solve(K_s, L)    # wrong implementation
        v = L.inverse() @ K_s   # correct implementation
        cov = K_ss - v.T @ v

        cov = cov.diag().view(-1, 1).expand_as(mu)
        return mu.squeeze(), cov
//  [python]
def choleskyinverse(upper, loss):
    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
    x_np = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]).astype(np.float32)
    x_np_zd = x_np.dot(x_np.T) + 1e-5 * np.eye(4)
    x_torch = torch.tensor(x_np_zd)
    torch_u = torch.cholesky(x_torch, upper)
    ms_u = Tensor(torch_u.numpy())
    choleskyinverse_ = CholeskyInverseNet(upper)
    choleskyinverse_output = choleskyinverse_(ms_u)
    choleskyinverse_expect = torch.cholesky_inverse(torch_u, upper)
    assert np.allclose(choleskyinverse_output.asnumpy(), choleskyinverse_expect.numpy(), loss, loss)
//  [python]
def choleskyinverse(upper, loss):
    context.set_context(mode=context.GRAPH_MODE, device_target='GPU')
    x_np = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]).astype(np.float32)
    x_np_zd = x_np.dot(x_np.T) + 1e-5 * np.eye(4)
    x_torch = torch.tensor(x_np_zd)
    torch_u = torch.linalg.cholesky(x_torch, upper=upper)
    ms_u = Tensor(torch_u.numpy())
    choleskyinverse_ = CholeskyInverseNet(upper)
    choleskyinverse_output = choleskyinverse_(ms_u)
    choleskyinverse_expect = torch.linalg.cholesky_inverse(torch_u, upper=upper)
    assert np.allclose(choleskyinverse_output.asnumpy(), choleskyinverse_expect.numpy(), loss, loss)
//  [python]
def forward(self, x1, x2, diag=False, **params):
        """
        Computes the Gaussian kernel matrix between inputs x1 and x2 belonging to a SPD manifold.

        Parameters
        ----------
        :param x1: input points on the SPD manifold
        :param x2: input points on the SPD manifold

        Optional parameters
        -------------------
        :param diag: Should we return the whole distance matrix, or just the diagonal? If True, we must have `x1 == x2`
        :param params: additional parameters

        Returns
        -------
        :return: kernel matrix between x1 and x2
        """
        # Transform input vector to matrix anc compute Cholesky decomposition
        x1 = torch.cholesky(vector_to_symmetric_matrix_mandel_torch(x1))
        x2 = torch.cholesky(vector_to_symmetric_matrix_mandel_torch(x2))

        # Expand dimensions to compute all matrix-matrix distances
        x1 = x1.unsqueeze(-3)
        x2 = x2.unsqueeze(-4)
        # Repeat x and y data along 1- and 0- dimensions to have ndata_x x ndata_y x dim x dim arrays
        x1 = torch.cat(x2.shape[-3] * [x1], dim=-3)
        x2 = torch.cat(x1.shape[-4] * [x2], dim=-4)

        # Compute singular values
        x1_inv_x2 = torch.bmm(x1.view(-1, self.dim, self.dim), torch.linalg.inv(x2).view(-1, self.dim, self.dim)).to(device)
        _, singular_values, _ = torch.linalg.svd(x1_inv_x2)

        # Reshape the singular values
        shape = list(x2.shape)[:-2]
        shape.append(singular_values.shape[-1])
        singular_values_rshp = singular_values.view(shape)

        # Compute difference and sum of squared singular values
        # Note: singular values that torch.linalg.svd outputs are sorted, the following code relies on this fact.
        H1, H2 = torch.log(singular_values_rshp[..., 0]), torch.log(singular_values_rshp[..., 1])
        # assert (H1 >= H2)
        r_H_sq = H1 ** 2 + H2 ** 2
        singular_values_diff = H1 - H2

        # Evaluate non-integral part
        nonintegral_part = torch.exp(-r_H_sq / (2 * self.lengthscale**2))

        # Evaluate integral part
        b_vals = torch.logspace(-5., 1., self.nb_points_integral).to(device)
        integral_vals = torch.zeros([self.nb_points_integral] + list(singular_values_diff.shape)).to(device)
        for i in range(self.nb_points_integral):
            integral_vals[i] = self.link_function(singular_values_diff, b_vals[i])
        # Kernel integral part
        kernel = torch.trapz(integral_vals, b_vals, dim=0) * nonintegral_part

        # Compute normalization factor
        # Evaluate only the integral part, as the non-integral part is equal to 1
        integral_vals_norm = torch.zeros([self.nb_points_integral]).to(device)
        for i in range(self.nb_points_integral):
            integral_vals_norm[i] = self.link_function(torch.zeros((1, 1), dtype=kernel.dtype).to(device), b_vals[i])
        # Kernel integral part
        norm_factor = torch.trapz(integral_vals_norm, b_vals, dim=0)

        # Kernel
        return kernel / norm_factor
//  [python]
def forward(self, x1, x2, diag=False, **params):
        """
        Computes the Gaussian kernel matrix between inputs x1 and x2 belonging to a SPD manifold.

        Parameters
        ----------
        :param x1: input points on the SPD manifold
        :param x2: input points on the SPD manifold

        Optional parameters
        -------------------
        :param diag: Should we return the whole distance matrix, or just the diagonal? If True, we must have `x1 == x2`
        :param params: additional parameters

        Returns
        -------
        :return: kernel matrix between x1 and x2
        """
        # Transform input vector to matrix anc compute Cholesky decomposition
        x1 = torch.linalg.cholesky(vector_to_symmetric_matrix_mandel_torch(x1))
        x2 = torch.linalg.cholesky(vector_to_symmetric_matrix_mandel_torch(x2))

        # Expand dimensions to compute all matrix-matrix distances
        x1 = x1.unsqueeze(-3)
        x2 = x2.unsqueeze(-4)
        # Repeat x and y data along 1- and 0- dimensions to have ndata_x x ndata_y x dim x dim arrays
        x1 = torch.cat(x2.shape[-3] * [x1], dim=-3)
        x2 = torch.cat(x1.shape[-4] * [x2], dim=-4)

        # Compute singular values
        x1_inv_x2 = torch.bmm(x1.view(-1, self.dim, self.dim), torch.linalg.inv(x2).view(-1, self.dim, self.dim)).to(device)
        _, singular_values, _ = torch.linalg.svd(x1_inv_x2)

        # Reshape the singular values
        shape = list(x2.shape)[:-2]
        shape.append(singular_values.shape[-1])
        singular_values_rshp = singular_values.view(shape)

        # Compute difference and sum of squared singular values
        # Note: singular values that torch.linalg.svd outputs are sorted, the following code relies on this fact.
        H1, H2 = torch.log(singular_values_rshp[..., 0]), torch.log(singular_values_rshp[..., 1])
        # assert (H1 >= H2)
        r_H_sq = H1 ** 2 + H2 ** 2
        singular_values_diff = H1 - H2

        # Evaluate non-integral part
        nonintegral_part = torch.exp(-r_H_sq / (2 * self.lengthscale**2))

        # Evaluate integral part
        b_vals = torch.logspace(-5., 1., self.nb_points_integral).to(device)
        integral_vals = torch.zeros([self.nb_points_integral] + list(singular_values_diff.shape)).to(device)
        for i in range(self.nb_points_integral):
            integral_vals[i] = self.link_function(singular_values_diff, b_vals[i])
        # Kernel integral part
        kernel = torch.trapz(integral_vals, b_vals, dim=0) * nonintegral_part

        # Compute normalization factor
        # Evaluate only the integral part, as the non-integral part is equal to 1
        integral_vals_norm = torch.zeros([self.nb_points_integral]).to(device)
        for i in range(self.nb_points_integral):
            integral_vals_norm[i] = self.link_function(torch.zeros((1, 1), dtype=kernel.dtype).to(device), b_vals[i])
        # Kernel integral part
        norm_factor = torch.trapz(integral_vals_norm, b_vals, dim=0)

        # Kernel
        return kernel / norm_factor
//  [python]
def choleskyinverse_pynative(upper, loss):
    context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
    x_np = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]).astype(np.float64)
    x_np_zd = x_np.dot(x_np.T) + 1e-5 * np.eye(4)
    x_torch = torch.tensor(x_np_zd)
    torch_u = torch.cholesky(x_torch, upper)
    ms_u = Tensor(torch_u.numpy())
    choleskyinverse_ = CholeskyInverseNet(upper)
    choleskyinverse_output = choleskyinverse_(ms_u)
    choleskyinverse_expect = torch.cholesky_inverse(torch_u, upper)
    assert np.allclose(choleskyinverse_output.asnumpy(), choleskyinverse_expect.numpy(), loss, loss)
//  [python]
def choleskyinverse_pynative(upper, loss):
    context.set_context(mode=context.PYNATIVE_MODE, device_target='GPU')
    x_np = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]]).astype(np.float64)
    x_np_zd = x_np.dot(x_np.T) + 1e-5 * np.eye(4)
    x_torch = torch.tensor(x_np_zd)
    torch_u = torch.linalg.cholesky(x_torch, upper=upper)
    ms_u = Tensor(torch_u.numpy())
    choleskyinverse_ = CholeskyInverseNet(upper)
    choleskyinverse_output = choleskyinverse_(ms_u)
    choleskyinverse_expect = torch.linalg.cholesky_inverse(torch_u, upper=upper)
    assert np.allclose(choleskyinverse_output.asnumpy(), choleskyinverse_expect.numpy(), loss, loss)
//  [python]
def cholesky_torch(S):
    # Cholesky decomposition
    # with adding a small value on the diagonal
    # until matrix is positive definite
    alpha = 0
    while 1:
        try:
            L = torch.cholesky(S + alpha*np.eye(*S.shape))
            return L
        except:
            if alpha == 0:
                alpha = 1e-10
            else:
                alpha *= 10
            print(">>>> {}::cholesky: Matrix is not positive definite, adding {:.0e} on the diagonal"
                .format(os.path.basename(__file__), alpha))
//  [python]
def cholesky_torch(S):
    # Cholesky decomposition
    # with adding a small value on the diagonal
    # until matrix is positive definite
    alpha = 0
    while 1:
        try:
            L = torch.linalg.cholesky(S + alpha*np.eye(*S.shape))
            return L
        except:
            if alpha == 0:
                alpha = 1e-10
            else:
                alpha *= 10
            print(">>>> {}::cholesky: Matrix is not positive definite, adding {:.0e} on the diagonal"
                .format(os.path.basename(__file__), alpha))
//  [python]
def set_covariance(self, k, cov):
        self.cov_factors[k].data = torch.cholesky(cov)
//  [python]
def set_covariance(self, k, cov):
        self.cov_factors[k].data = torch.linalg.cholesky(cov)
//  [python]
def convert_cbc_terms_to_socp_terms(bfe, e, V, bfv, v, extravars,
                                        testing=False):
        m = bfe.shape[-1]
        bfc = bfe.new_zeros((m+extravars))
        if testing:
            u0 = torch.zeros((m,))
            u0_hom = torch.cat((torch.tensor([1.]), u0))
        with torch.no_grad():
            # [1, u] Asq [1; u]
            Asq = torch.cat(
                (
                    torch.cat((torch.tensor([[v]]),         (bfv / 2).reshape(1, -1)), dim=-1),
                    torch.cat(((bfv / 2).reshape(-1, 1),    V), dim=-1)
                ),
                dim=-2)
            if testing:
                np.testing.assert_allclose(u0_hom @ Asq @ u0_hom,
                                        u0 @ V @ u0 + bfv @ u0 + v)

            # [1, u] Asq [1; u] = |L[1; u]|_2 = |[0, A] [y_1; y_2; u] + b|_2
            n_constraints = m+1
            try:
                L = torch.cholesky(Asq) # (m+1) x (m+1)
            except RuntimeError as err:
                if "cholesky" in str(err) and "singular" in str(err):
                    L = torch.cholesky(Asq + torch.diag(torch.ones(m+1)*1e-3))
                else:
                    raise
            # try:
            #     n_constraints = m+1
            #     L = torch.cholesky(Asq) # (m+1) x (m+1)
            # except RuntimeError as err:
            #     if "cholesky" in str(err) and "singular" in str(err):
            #         if torch.allclose(v, torch.zeros((1,))) and torch.allclose(bfv, torch.zeros(bfv.shape[0])):
            #             L = torch.zeros((m+1, m+1))
            #             L[1:, 1:] = torch.cholesky(V)
            #         else:
            #             diag_e, U = torch.symeig(Asq, eigenvectors=True)
            #             n_constraints = torch.nonzero(diag_e, as_tuple=True)[0].shape[-1]
            #             L = torch.diag(diag_e).sqrt() @ U[:, -n_constraints:]
            #     else:
            #         raise

        if testing:
            np.testing.assert_allclose(L @ L.T, Asq, rtol=1e-2, atol=1e-3)

        A = torch.zeros((n_constraints, m + extravars))
        A[:, extravars:] = L.T[:, 1:]
        bfb = L.T[:, 0] # (m+1)
        if testing:
            y_u0 = torch.cat((torch.zeros(extravars), u0))
            np.testing.assert_allclose(A @ y_u0 + bfb, L.T @ u0_hom)
            np.testing.assert_allclose(u0_hom @ Asq @ u0_hom, u0_hom @ L @ L.T @ u0_hom, rtol=1e-2, atol=1e-3)
            np.testing.assert_allclose(u0 @ V @ u0 + bfv @ u0 + v, (A @ y_u0 + bfb) @ (A @ y_u0 + bfb), rtol=1e-2, atol=1e-3)
        assert extravars >= 1, "I assumed atleast δ "
        bfc[extravars-1] = 1 # For delta the relaxation factor
        bfc[extravars:] = bfe # only affine terms need to be negative?
        d = e
        return A, bfb, bfc, d
//  [python]
def convert_cbc_terms_to_socp_terms(bfe, e, V, bfv, v, extravars,
                                        testing=False):
        m = bfe.shape[-1]
        bfc = bfe.new_zeros((m+extravars))
        if testing:
            u0 = torch.zeros((m,))
            u0_hom = torch.cat((torch.tensor([1.]), u0))
        with torch.no_grad():
            # [1, u] Asq [1; u]
            Asq = torch.cat(
                (
                    torch.cat((torch.tensor([[v]]),         (bfv / 2).reshape(1, -1)), dim=-1),
                    torch.cat(((bfv / 2).reshape(-1, 1),    V), dim=-1)
                ),
                dim=-2)
            if testing:
                np.testing.assert_allclose(u0_hom @ Asq @ u0_hom,
                                        u0 @ V @ u0 + bfv @ u0 + v)

            # [1, u] Asq [1; u] = |L[1; u]|_2 = |[0, A] [y_1; y_2; u] + b|_2
            n_constraints = m+1
            try:
                L = torch.linalg.cholesky(Asq) # (m+1) x (m+1)
            except RuntimeError as err:
                if "cholesky" in str(err) and "singular" in str(err):
                    L = torch.linalg.cholesky(Asq + torch.diag(torch.ones(m+1)*1e-3))
                else:
                    raise
            # try:
            #     n_constraints = m+1
            #     L = torch.linalg.cholesky(Asq) # (m+1) x (m+1)
            # except RuntimeError as err:
            #     if "cholesky" in str(err) and "singular" in str(err):
            #         if torch.allclose(v, torch.zeros((1,))) and torch.allclose(bfv, torch.zeros(bfv.shape[0])):
            #             L = torch.zeros((m+1, m+1))
            #             L[1:, 1:] = torch.linalg.cholesky(V)
            #         else:
            #             diag_e, U = torch.symeig(Asq, eigenvectors=True)
            #             n_constraints = torch.nonzero(diag_e, as_tuple=True)[0].shape[-1]
            #             L = torch.diag(diag_e).sqrt() @ U[:, -n_constraints:]
            #     else:
            #         raise

        if testing:
            np.testing.assert_allclose(L @ L.T, Asq, rtol=1e-2, atol=1e-3)

        A = torch.zeros((n_constraints, m + extravars))
        A[:, extravars:] = L.T[:, 1:]
        bfb = L.T[:, 0] # (m+1)
        if testing:
            y_u0 = torch.cat((torch.zeros(extravars), u0))
            np.testing.assert_allclose(A @ y_u0 + bfb, L.T @ u0_hom)
            np.testing.assert_allclose(u0_hom @ Asq @ u0_hom, u0_hom @ L @ L.T @ u0_hom, rtol=1e-2, atol=1e-3)
            np.testing.assert_allclose(u0 @ V @ u0 + bfv @ u0 + v, (A @ y_u0 + bfb) @ (A @ y_u0 + bfb), rtol=1e-2, atol=1e-3)
        assert extravars >= 1, "I assumed atleast δ "
        bfc[extravars-1] = 1 # For delta the relaxation factor
        bfc[extravars:] = bfe # only affine terms need to be negative?
        d = e
        return A, bfb, bfc, d
//  [python]
def train(self, K, y, method='cholesky', alpha=1e-3):
        """
        Compute the output weights with a linear regression.
        Parameters:
          - K: input sequence of shape (n_input, n_input)
          - y: target output (n_input, out_dim)
          - method: "cholesky" or "sklearn ridge"
          - alpha: L2-regularization parameter
        
        Returns: a tensor of shape (n_input, out_dim)
        """
        if method == 'cholesky':
            # This technique uses the Cholesky decomposition
            K.view(-1)[::len(K)+1] += alpha  # add elements on the diagonal inplace
            L = torch.cholesky(K, upper=False)
            return torch.cholesky_solve(y, L, upper=False)
        elif method == 'sklearn ridge':
            from sklearn.linear_model import Ridge
            clf = Ridge(fit_intercept=False, alpha=alpha)
            clf.fit(K.cpu().numpy(), y.cpu().numpy())
            return torch.from_numpy(clf.coef_.T).to(self.device)
//  [python]
def train(self, K, y, method='cholesky', alpha=1e-3):
        """
        Compute the output weights with a linear regression.
        Parameters:
          - K: input sequence of shape (n_input, n_input)
          - y: target output (n_input, out_dim)
          - method: "cholesky" or "sklearn ridge"
          - alpha: L2-regularization parameter
        
        Returns: a tensor of shape (n_input, out_dim)
        """
        if method == 'cholesky':
            # This technique uses the Cholesky decomposition
            K.view(-1)[::len(K)+1] += alpha  # add elements on the diagonal inplace
            L = torch.linalg.cholesky(K, upper=False)
            return torch.linalg.solve_triangular(L, y, upper=False)
        elif method == 'sklearn ridge':
            from sklearn.linear_model import Ridge
            clf = Ridge(fit_intercept=False, alpha=alpha)
            clf.fit(K.cpu().numpy(), y.cpu().numpy())
            return torch.from_numpy(clf.coef_.T).to(self.device)
//  [python]
def log_prob(self, x):
        """
        x: N x n_features
        returns: N x n_states
        """
        identities = torch.eye(self.n_feats).repeat(self.n_states, 1, 1).to(self.device)
        L = torch.cholesky(self.sigma, upper=False)
        L_inv, _ = torch.triangular_solve(L, identities, upper=False)
        precisions = L_inv.contiguous()  # Inverse sigma

        mu = self.mu

        exp_term = torch.zeros(x.size(0), self.n_states).to(self.device)

        for k in range(self.n_states):
            y = torch.mm(x, precisions[k]) - torch.mv(precisions[k], mu[k])
            exp_term[:, k] = torch.sum(y * y, dim=1)

        log_det = torch.sum(
            torch.log(precisions.view(self.n_states, -1)[:, :: self.n_feats + 1]),
            dim=1,
        )

        logp = -0.5 * (self.n_feats * math.log(2 * math.pi) + exp_term) + log_det

        return logp
//  [python]
def log_prob(self, x):
        """
        x: N x n_features
        returns: N x n_states
        """
        identities = torch.eye(self.n_feats).repeat(self.n_states, 1, 1).to(self.device)
        L = torch.linalg.cholesky(self.sigma, upper=False)
        L_inv, _ = torch.triangular_solve(L, identities, upper=False)
        precisions = L_inv.contiguous()  # Inverse sigma

        mu = self.mu

        exp_term = torch.zeros(x.size(0), self.n_states).to(self.device)

        for k in range(self.n_states):
            y = torch.mm(x, precisions[k]) - torch.mv(precisions[k], mu[k])
            exp_term[:, k] = torch.sum(y * y, dim=1)

        log_det = torch.sum(
            torch.log(precisions.view(self.n_states, -1)[:, :: self.n_feats + 1]),
            dim=1,
        )

        logp = -0.5 * (self.n_feats * math.log(2 * math.pi) + exp_term) + log_det

        return logp
//  [python]
def update(self, x, P, z, mask, idx):
        """
        Update state x and P after the observation, outputs filtered state and covariance
        Masking done according to formulas from
        https://www.stat.pitt.edu/stoffer/tsa4/tsa4.pdf, Chapter 6.4 Page 311
        https://www.stat.pitt.edu/stoffer/dss_files/em.pdf, Chapter 3
        """
        if self.fixed_params:
            # create one parameter matrix for each observed sample in batch
            pre_R_masked = self.pre_R.clone()
            pre_R_masked = pre_R_masked.repeat(mask.shape[0], 1)
            identities = torch.eye(self.dim, self.dim, device=self.device).repeat(mask.shape[0], 1, 1)
            H_masked = self.H.clone().repeat(mask.shape[0], 1, 1)    
        else:
            # create one parameter matrix for each observed sample in batch
            pre_R_masked = self.pre_R[idx].clone().squeeze(1)
            identities = torch.eye(self.dim, self.dim, device=self.device).repeat(mask.shape[0], 1, 1)
            H_masked = self.H[idx].clone()

        # mask parameter matrices for partial measurements
        pre_R_masked[~mask.to(torch.bool)] = 1
        R_masked = identities * (pre_R_masked**2).unsqueeze(-1)
        H_masked[~mask.to(torch.bool)] = 0
        # masked emission
        masked_sigma = H_masked @ P @ H_masked.transpose(1, 2) + R_masked
        # Kalman gain, a more stable implementation than naive P @ H^T @ y_sigma^{-1}
        L = torch.cholesky(masked_sigma)
        K = torch.triangular_solve(H_masked @ P.transpose(1, 2), L, upper=False)[0]
        K = torch.triangular_solve(K, L.transpose(1, 2))[0].transpose(1, 2)
        z = (self.flow.inverse(z.squeeze(-1))[0] * mask).unsqueeze(-1)
        # Update state mean and covariance p(x | y)
        v = z - H_masked @ x
        x = x + K @ v  
        identity = torch.eye(*P.shape[1:], device=torch.device(self.device))
        # Joseph Form for numerical stability
        P = (identity - K @ H_masked) @ P @ (identity - K @ H_masked).transpose(1, 2) + K @ R_masked @ K.transpose(1, 2)

        return x, P
//  [python]
def update(self, x, P, z, mask, idx):
        """
        Update state x and P after the observation, outputs filtered state and covariance
        Masking done according to formulas from
        https://www.stat.pitt.edu/stoffer/tsa4/tsa4.pdf, Chapter 6.4 Page 311
        https://www.stat.pitt.edu/stoffer/dss_files/em.pdf, Chapter 3
        """
        if self.fixed_params:
            # create one parameter matrix for each observed sample in batch
            pre_R_masked = self.pre_R.clone()
            pre_R_masked = pre_R_masked.repeat(mask.shape[0], 1)
            identities = torch.eye(self.dim, self.dim, device=self.device).repeat(mask.shape[0], 1, 1)
            H_masked = self.H.clone().repeat(mask.shape[0], 1, 1)    
        else:
            # create one parameter matrix for each observed sample in batch
            pre_R_masked = self.pre_R[idx].clone().squeeze(1)
            identities = torch.eye(self.dim, self.dim, device=self.device).repeat(mask.shape[0], 1, 1)
            H_masked = self.H[idx].clone()

        # mask parameter matrices for partial measurements
        pre_R_masked[~mask.to(torch.bool)] = 1
        R_masked = identities * (pre_R_masked**2).unsqueeze(-1)
        H_masked[~mask.to(torch.bool)] = 0
        # masked emission
        masked_sigma = H_masked @ P @ H_masked.transpose(1, 2) + R_masked
        # Kalman gain, a more stable implementation than naive P @ H^T @ y_sigma^{-1}
        L = torch.linalg.cholesky(masked_sigma)
        K = torch.triangular_solve(H_masked @ P.transpose(1, 2), L, upper=False)[0]
        K = torch.triangular_solve(K, L.transpose(1, 2))[0].transpose(1, 2)
        z = (self.flow.inverse(z.squeeze(-1))[0] * mask).unsqueeze(-1)
        # Update state mean and covariance p(x | y)
        v = z - H_masked @ x
        x = x + K @ v  
        identity = torch.eye(*P.shape[1:], device=torch.device(self.device))
        # Joseph Form for numerical stability
        P = (identity - K @ H_masked) @ P @ (identity - K @ H_masked).transpose(1, 2) + K @ R_masked @ K.transpose(1, 2)

        return x, P
//  [python]
def cholesky_solve_test_helper(self, A_dims, b_dims, upper, device, dtype):
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix

        b = torch.randn(*b_dims, dtype=dtype, device=device)
        A = random_symmetric_pd_matrix(*A_dims, dtype=dtype, device=device)
        L = torch.cholesky(A, upper=upper)
        return b, A, L
//  [python]
def cholesky_solve_test_helper(self, A_dims, b_dims, upper, device, dtype):
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix

        b = torch.randn(*b_dims, dtype=dtype, device=device)
        A = random_symmetric_pd_matrix(*A_dims, dtype=dtype, device=device)
        L = torch.linalg.cholesky(A, upper=upper)
        return b, A, L
//  [python]
def __init__(self, k=None, dim=None, means=None, covariances=None, weights=None, normalize=False):
        super(GMM, self).__init__()
        if k is None and means is None:
            raise NotImplementedError

        if means is None:
            covars = torch.rand(k, dim, dim)
            covars = torch.matmul(covars, covars.transpose(1, 2))
            self.means = nn.ParameterList([nn.Parameter(m) for m in torch.randn(k, dim)])
            self.cov_factors = nn.ParameterList([nn.Parameter(torch.cholesky(cov)) for cov in covars])
            self.weights = nn.Parameter(torch.FloatTensor([1./k] * k))
            self.k = k
        else:
            self.means = nn.ParameterList([nn.Parameter(m) for m in means])
            self.cov_factors = nn.ParameterList([nn.Parameter(torch.cholesky(cov)) for cov in covariances])
            self.weights = nn.Parameter(weights)
            self.k = weights.shape[0]

        self.normalize = normalize
//  [python]
def __init__(self, k=None, dim=None, means=None, covariances=None, weights=None, normalize=False):
        super(GMM, self).__init__()
        if k is None and means is None:
            raise NotImplementedError

        if means is None:
            covars = torch.rand(k, dim, dim)
            covars = torch.matmul(covars, covars.transpose(1, 2))
            self.means = nn.ParameterList([nn.Parameter(m) for m in torch.randn(k, dim)])
            self.cov_factors = nn.ParameterList([nn.Parameter(torch.linalg.cholesky(cov)) for cov in covars])
            self.weights = nn.Parameter(torch.FloatTensor([1./k] * k))
            self.k = k
        else:
            self.means = nn.ParameterList([nn.Parameter(m) for m in means])
            self.cov_factors = nn.ParameterList([nn.Parameter(torch.linalg.cholesky(cov)) for cov in covariances])
            self.weights = nn.Parameter(weights)
            self.k = weights.shape[0]

        self.normalize = normalize
//  [python]
def test_cholesky_batched(self, device, dtype):
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix

        def cholesky_test_helper(n, batch_dims, upper):
            A = random_symmetric_pd_matrix(n, *batch_dims, dtype=dtype, device=device)
            cholesky_exp = torch.stack([m.cholesky(upper=upper) for m in A.reshape(-1, n, n)])
            cholesky_exp = cholesky_exp.reshape_as(A)
            self.assertEqual(cholesky_exp, torch.cholesky(A, upper=upper))

        for upper, batchsize in product([True, False], [(3,), (3, 4), (2, 3, 4)]):
            cholesky_test_helper(3, batchsize, upper)
//  [python]
def test_cholesky_batched(self, device, dtype):
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix

        def cholesky_test_helper(n, batch_dims, upper):
            A = random_symmetric_pd_matrix(n, *batch_dims, dtype=dtype, device=device)
            cholesky_exp = torch.stack([torch.linalg.cholesky(m, upper=upper) for m in A.reshape(-1, n, n)])
            cholesky_exp = cholesky_exp.reshape_as(A)
            self.assertEqual(cholesky_exp, torch.linalg.cholesky(A, upper=upper))

        for upper, batchsize in product([True, False], [(3,), (3, 4), (2, 3, 4)]):
            cholesky_test_helper(3, batchsize, upper)
//  [python]
def test_cholesky_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix

        for upper in [True, False]:
            A = random_symmetric_pd_matrix(2, 2, dtype=dtype, device='cpu')
            b = torch.randn(2, 2, 2, dtype=dtype, device='cpu')
            x_exp = torch.Tensor(solve(A.permute(0, 2, 1).numpy(), b.permute(2, 1, 0).numpy())).to(dtype=dtype, device=device)
            A = A.to(device).permute(0, 2, 1)
            b = b.to(device).permute(2, 1, 0)
            assert not A.is_contiguous() and not b.is_contiguous(), "contiguous inputs"
            L = torch.cholesky(A, upper)
            x = torch.cholesky_solve(b, L, upper=upper)
            self.assertEqual(x, x_exp)
//  [python]
def test_cholesky_solve_batched_non_contiguous(self, device, dtype):
        from numpy.linalg import solve
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix

        for upper in [True, False]:
            A = random_symmetric_pd_matrix(2, 2, dtype=dtype, device='cpu')
            b = torch.randn(2, 2, 2, dtype=dtype, device='cpu')
            x_exp = torch.Tensor(solve(A.permute(0, 2, 1).numpy(), b.permute(2, 1, 0).numpy())).to(dtype=dtype, device=device)
            A = A.to(device).permute(0, 2, 1)
            b = b.to(device).permute(2, 1, 0)
            assert not A.is_contiguous() and not b.is_contiguous(), "contiguous inputs"
            L = torch.linalg.cholesky(A, upper)
            x = torch.cholesky_solve(b, L, upper=upper)
            self.assertEqual(x, x_exp)
//  [python]
def test_cholesky(self, device, dtype):
        x = torch.rand(10, 10, dtype=dtype, device=device) + 1e-1
        A = torch.mm(x, x.t())

        # default Case
        C = torch.cholesky(A)
        B = torch.mm(C, C.t())
        self.assertEqual(A, B, atol=1e-14, rtol=0)

        # test Upper Triangular
        U = torch.cholesky(A, True)
        B = torch.mm(U.t(), U)
        self.assertEqual(A, B, atol=1e-14, rtol=0, msg='cholesky (upper) did not allow rebuilding the original matrix')

        # test Lower Triangular
        L = torch.cholesky(A, False)
        B = torch.mm(L, L.t())
        self.assertEqual(A, B, atol=1e-14, rtol=0, msg='cholesky (lower) did not allow rebuilding the original matrix')
//  [python]
def test_cholesky(self, device, dtype):
        x = torch.rand(10, 10, dtype=dtype, device=device) + 1e-1
        A = torch.mm(x, x.t())

        # default Case
        C = torch.linalg.cholesky(A)
        B = torch.mm(C, C.t())
        self.assertEqual(A, B, atol=1e-14, rtol=0)

        # test Upper Triangular
        U = torch.linalg.cholesky(A, upper=True)
        B = torch.mm(U.t(), U)
        self.assertEqual(A, B, atol=1e-14, rtol=0, msg='cholesky (upper) did not allow rebuilding the original matrix')

        # test Lower Triangular
        L = torch.linalg.cholesky(A, upper=False)
        B = torch.mm(L, L.t())
        self.assertEqual(A, B, atol=1e-14, rtol=0, msg='cholesky (lower) did not allow rebuilding the original matrix')
//  [python]
def test_cholesky_batched_many_batches(self, device, dtype):
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix

        def cholesky_test_helper(n, batchsize, device, upper):
            A = random_symmetric_pd_matrix(n, batchsize, dtype=dtype, device=device)
            chol_fact = torch.cholesky(A, upper=upper)
            if upper:
                # Correctness check
                self.assertEqual(A, chol_fact.transpose(-2, -1).matmul(chol_fact))
                # Upper triangular check
                self.assertEqual(chol_fact, chol_fact.triu())
            else:
                # Correctness check
                self.assertEqual(A, chol_fact.matmul(chol_fact.transpose(-2, -1)))
                # Lower triangular check
                self.assertEqual(chol_fact, chol_fact.tril())

        for upper, batchsize in product([True, False], [262144, 524288]):
            cholesky_test_helper(2, batchsize, device, upper)
//  [python]
def test_cholesky_batched_many_batches(self, device, dtype):
        from torch.testing._internal.common_utils import random_symmetric_pd_matrix

        def cholesky_test_helper(n, batchsize, device, upper):
            A = random_symmetric_pd_matrix(n, batchsize, dtype=dtype, device=device)
            chol_fact = torch.linalg.cholesky(A, upper=upper)
            if upper:
                # Correctness check
                self.assertEqual(A, chol_fact.transpose(-2, -1).matmul(chol_fact))
                # Upper triangular check
                self.assertEqual(chol_fact, chol_fact.triu())
            else:
                # Correctness check
                self.assertEqual(A, chol_fact.matmul(chol_fact.transpose(-2, -1)))
                # Lower triangular check
                self.assertEqual(chol_fact, chol_fact.tril())

        for upper, batchsize in product([True, False], [262144, 524288]):
            cholesky_test_helper(2, batchsize, device, upper)