jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def forward_routine(layout, device, value_dtype, index_dtype, shapes, upper, unitriangular, transpose):
    if index_dtype == torch.int32 and layout is torch.sparse_coo:
        pytest.skip("Skipping test as sparse COO tensors with int32 indices are not supported")
    if platform == "win32" and device == torch.device("cpu"):
        pytest.skip("Skipping triangular solve CPU tests as solver not implemented for Windows OS")

    _, A_shape, B_shape, A_nnz = shapes
    A = rand_sparse_tri(
        A_shape,
        A_nnz,
        layout,
        upper=upper,
        strict=unitriangular,
        indices_dtype=index_dtype,
        values_dtype=value_dtype,
        device=device,
    )
    B = torch.rand(*B_shape, dtype=value_dtype, device=device)
    Ad = A.to_dense()

    res_ref = torch.triangular_solve(B, Ad, upper=upper, unitriangular=unitriangular, transpose=transpose).solution
    res_test = sparse_triangular_solve(A, B, upper=upper, unitriangular=unitriangular, transpose=transpose)

    assert torch.allclose(res_test, res_ref, atol=ATOL, rtol=RTOL)
//  [python]
def forward_routine(layout, device, value_dtype, index_dtype, shapes, upper, unitriangular, transpose):
    if index_dtype == torch.int32 and layout is torch.sparse_coo:
        pytest.skip("Skipping test as sparse COO tensors with int32 indices are not supported")
    if platform == "win32" and device == torch.device("cpu"):
        pytest.skip("Skipping triangular solve CPU tests as solver not implemented for Windows OS")

    _, A_shape, B_shape, A_nnz = shapes
    A = rand_sparse_tri(
        A_shape,
        A_nnz,
        layout,
        upper=upper,
        strict=unitriangular,
        indices_dtype=index_dtype,
        values_dtype=value_dtype,
        device=device,
    )
    B = torch.rand(*B_shape, dtype=value_dtype, device=device)
    Ad = A.to_dense()

    res_ref = torch.linalg.solve_triangular(Ad, B, upper=upper, unitriangular=unitriangular, transpose=transpose)
    res_test = sparse_triangular_solve(A, B, upper=upper, unitriangular=unitriangular, transpose=transpose)

    assert torch.allclose(res_test, res_ref, atol=ATOL, rtol=RTOL)
//  [python]
def triangular_solve(b, A):
    s, c = torch.triangular_solve(b, A)
    return s
//  [python]
def triangular_solve(b, A):
    s, c = torch.linalg.solve_triangular(b, A)
    return s
//  [python]
def sample(self, num_samples, context):
        """
        Generated num_samples independent samples from p(inputs | context).
        NB: Generates num_samples samples for EACH item in context batch i.e. returns
        (num_samples * batch_size) samples in total.

        :param num_samples: int
            Number of samples to generate.
        :param context: torch.Tensor [batch_size, context_dim]
            Conditioning variable.
        :return: torch.Tensor [batch_size, num_samples, output_dim]
            Batch of generated samples.
        """

        # Get necessary quantities.
        logits, means, _, _, precision_factors = self.get_mixture_components(context)
        batch_size, n_mixtures, output_dim = means.shape

        # We need (batch_size * num_samples) samples in total.
        means, precision_factors = (
            utils.repeat_rows(means, num_samples),
            utils.repeat_rows(precision_factors, num_samples),
        )

        # Normalize the logits for the coefficients.
        coefficients = F.softmax(logits, dim=-1)  # [batch_size, num_components]

        # Choose num_samples mixture components per example in the batch.
        choices = torch.multinomial(
            coefficients, num_samples=num_samples, replacement=True
        ).view(
            -1
        )  # [batch_size, num_samples]

        # Create dummy index for indexing means and precision factors.
        ix = utils.repeat_rows(torch.arange(batch_size), num_samples)

        # Select means and precision factors.
        chosen_means = means[ix, choices, :]
        chosen_precision_factors = precision_factors[ix, choices, :, :]

        # Batch triangular solve to multiply standard normal samples by inverse
        # of upper triangular precision factor.
        zero_mean_samples, _ = torch.triangular_solve(
            torch.randn(
                batch_size * num_samples, output_dim, 1
            ),  # Need dummy final dimension.
            chosen_precision_factors,
        )

        # Mow center samples at chosen means, removing dummy final dimension
        # from triangular solve.
        samples = chosen_means + zero_mean_samples.squeeze(-1)

        return samples.reshape(batch_size, num_samples, output_dim)
//  [python]
def sample(self, num_samples, context):
        """
        Generated num_samples independent samples from p(inputs | context).
        NB: Generates num_samples samples for EACH item in context batch i.e. returns
        (num_samples * batch_size) samples in total.

        :param num_samples: int
            Number of samples to generate.
        :param context: torch.Tensor [batch_size, context_dim]
            Conditioning variable.
        :return: torch.Tensor [batch_size, num_samples, output_dim]
            Batch of generated samples.
        """

        # Get necessary quantities.
        logits, means, _, _, precision_factors = self.get_mixture_components(context)
        batch_size, n_mixtures, output_dim = means.shape

        # We need (batch_size * num_samples) samples in total.
        means, precision_factors = (
            utils.repeat_rows(means, num_samples),
            utils.repeat_rows(precision_factors, num_samples),
        )

        # Normalize the logits for the coefficients.
        coefficients = F.softmax(logits, dim=-1)  # [batch_size, num_components]

        # Choose num_samples mixture components per example in the batch.
        choices = torch.multinomial(
            coefficients, num_samples=num_samples, replacement=True
        ).view(
            -1
        )  # [batch_size, num_samples]

        # Create dummy index for indexing means and precision factors.
        ix = utils.repeat_rows(torch.arange(batch_size), num_samples)

        # Select means and precision factors.
        chosen_means = means[ix, choices, :]
        chosen_precision_factors = precision_factors[ix, choices, :, :]

        # Batch triangular solve to multiply standard normal samples by inverse
        # of upper triangular precision factor.
        zero_mean_samples, _ = torch.linalg.solve_triangular(
            torch.randn(
                batch_size * num_samples, output_dim, 1
            ),  # Need dummy final dimension.
            chosen_precision_factors,
        )

        # Mow center samples at chosen means, removing dummy final dimension
        # from triangular solve.
        samples = chosen_means + zero_mean_samples.squeeze(-1)

        return samples.reshape(batch_size, num_samples, output_dim)
//  [python]
def inv_matmul(self, right_tensor: Tensor, left_tensor: Optional[Tensor] = None) -> Tensor:
        if isinstance(self._tensor, DenseLinearOperator):
            res = torch.triangular_solve(right_tensor, self.to_dense(), upper=self.upper).solution
        elif isinstance(self._tensor, BatchRepeatLinearOperator):
            res = self._tensor.base_linear_operator.inv_matmul(right_tensor, left_tensor)
            # TODO: Proper broadcasting
            res = res.expand(self._tensor.batch_repeat + res.shape[-2:])
        else:
            # TODO: Can we be smarter here?
            res = self._tensor.inv_matmul(right_tensor=right_tensor, left_tensor=left_tensor)
        if left_tensor is not None:
            res = left_tensor @ res
        return res
//  [python]
def inv_matmul(self, right_tensor: Tensor, left_tensor: Optional[Tensor] = None) -> Tensor:
        if isinstance(self._tensor, DenseLinearOperator):
            res = torch.linalg.solve_triangular(self.to_dense(), right_tensor, upper=self.upper)
        elif isinstance(self._tensor, BatchRepeatLinearOperator):
            res = self._tensor.base_linear_operator.inv_matmul(right_tensor, left_tensor)
            # TODO: Proper broadcasting
            res = res.expand(self._tensor.batch_repeat + res.shape[-2:])
        else:
            # TODO: Can we be smarter here?
            res = self._tensor.inv_matmul(right_tensor=right_tensor, left_tensor=left_tensor)
        if left_tensor is not None:
            res = left_tensor @ res
        return res
//  [python]
def forward(self, X, Y):
        self.x_train, self.y_train = X.detach(), Y.detach()

        L, A, LB, d = self._calc_intermediate(X, Y)
        tmp3 = Y - A.T @ d
        numerator = torch.square(tmp3)
        # Denominator (1/n * Tr(I - H))^2
        C = torch.triangular_solve(A, LB, upper=False)[0]
        denominator = (1 - torch.square(C).sum() / X.shape[0]) ** 2
        loss = (numerator / denominator).mean()

        self._save_losses(loss)
        return loss
//  [python]
def forward(self, X, Y):
        self.x_train, self.y_train = X.detach(), Y.detach()

        L, A, LB, d = self._calc_intermediate(X, Y)
        tmp3 = Y - A.T @ d
        numerator = torch.square(tmp3)
        # Denominator (1/n * Tr(I - H))^2
        C = torch.linalg.solve_triangular(LB, A, upper=False)
        denominator = (1 - torch.square(C).sum() / X.shape[0]) ** 2
        loss = (numerator / denominator).mean()

        self._save_losses(loss)
        return loss
//  [python]
def predict(self, X):
        if self.x_train is None or self.y_train is None:
            raise RuntimeError("Call forward at least once before calling predict.")
        with torch.autograd.no_grad():
            L, A, LB, d = self._calc_intermediate(self.x_train, self.y_train)
            sqrt_var = torch.sqrt(self.penalty * (X.shape[0] - 1))
            kms = self.kernel(self.centers, X)
            tmp1 = torch.triangular_solve(d / sqrt_var, L.T, upper=True)[0]
            return kms.T @ tmp1
//  [python]
def predict(self, X):
        if self.x_train is None or self.y_train is None:
            raise RuntimeError("Call forward at least once before calling predict.")
        with torch.autograd.no_grad():
            L, A, LB, d = self._calc_intermediate(self.x_train, self.y_train)
            sqrt_var = torch.sqrt(self.penalty * (X.shape[0] - 1))
            kms = self.kernel(self.centers, X)
            tmp1 = torch.linalg.solve_triangular(L.T, d / sqrt_var, upper=True)
            return kms.T @ tmp1
//  [python]
def get_white_from_latents(latents, const_mean, cholesky_gram):
    """
    latents: (B,D)
    const_mean: scalar
    cholesky_gram: (D,D)
    """
    white = torch.triangular_solve(
        latents.unsqueeze(-1) - const_mean, cholesky_gram, upper=False
    )[0]
    return white.squeeze(dim=-1)
//  [python]
def get_white_from_latents(latents, const_mean, cholesky_gram):
    """
    latents: (B,D)
    const_mean: scalar
    cholesky_gram: (D,D)
    """
    white = torch.linalg.solve_triangular(
        cholesky_gram, latents.unsqueeze(-1) - const_mean, upper=False
    )
    return white.squeeze(dim=-1)
//  [python]
def _standardize(self, x: Tensor) -> Tensor:
        """_transform helper standardizing the input `x`, which is assumed to be a
        `d x 1`-dim Tensor, or a `len(self._not_fixed) x 1`-dim if there are fixed
        indices.
        """
        z = x
        if self._mean is not None:
            z = z - self._mean
        root = self._covariance_root
        if root is not None:
            z = torch.triangular_solve(z, root, upper=False)[0]

        return z
//  [python]
def _standardize(self, x: Tensor) -> Tensor:
        """_transform helper standardizing the input `x`, which is assumed to be a
        `d x 1`-dim Tensor, or a `len(self._not_fixed) x 1`-dim if there are fixed
        indices.
        """
        z = x
        if self._mean is not None:
            z = z - self._mean
        root = self._covariance_root
        if root is not None:
            z = torch.linalg.solve_triangular(root, z, upper=False)

        return z
//  [python]
def _calc_intermediate(self, X, Y):
        # Like with LOOCV we are virtually using an estimator trained with n - 1 points.
        variance = self.penalty * (X.shape[0] - 1)
        sqrt_var = torch.sqrt(variance)

        kmn = self.kernel(self.centers, X)
        kmm = self.kernel(self.centers, self.centers)
        L = jittering_cholesky(kmm)  # L @ L.T = kmm
        # A = L^{-1} K_mn / (sqrt(n*pen))
        A = torch.triangular_solve(kmn, L, upper=False).solution / sqrt_var
        AAT = A @ A.T
        # B = A @ A.T + I
        B = AAT + torch.eye(AAT.shape[0], device=X.device, dtype=X.dtype)
        LB = jittering_cholesky(B)  # LB @ LB.T = B

        AY = A @ Y
        # numerator is (1/n)*||(I - A.T @ LB^{-T} @ LB^{-1} @ A) @ Y||^2
        # compute A.T @ LB^{-T} @ LB^{-1} @ A @ Y
        tmp1 = torch.triangular_solve(AY, LB, upper=False).solution
        tmp2 = torch.triangular_solve(tmp1, LB.T, upper=True).solution
        d = tmp2
        return L, A, LB, d
//  [python]
def _calc_intermediate(self, X, Y):
        # Like with LOOCV we are virtually using an estimator trained with n - 1 points.
        variance = self.penalty * (X.shape[0] - 1)
        sqrt_var = torch.sqrt(variance)

        kmn = self.kernel(self.centers, X)
        kmm = self.kernel(self.centers, self.centers)
        L = jittering_cholesky(kmm)  # L @ L.T = kmm
        # A = L^{-1} K_mn / (sqrt(n*pen))
        A = torch.linalg.solve_triangular(L, kmn, upper=False) / sqrt_var
        AAT = A @ A.T
        # B = A @ A.T + I
        B = AAT + torch.eye(AAT.shape[0], device=X.device, dtype=X.dtype)
        LB = jittering_cholesky(B)  # LB @ LB.T = B

        AY = A @ Y
        # numerator is (1/n)*||(I - A.T @ LB^{-T} @ LB^{-1} @ A) @ Y||^2
        # compute A.T @ LB^{-T} @ LB^{-1} @ A @ Y
        tmp1 = torch.linalg.solve_triangular(LB, AY, upper=False)
        tmp2 = torch.linalg.solve_triangular(LB.T, tmp1, upper=True)
        d = tmp2
        return L, A, LB, d
//  [python]
def calc_trace_fwd(init_val: torch.Tensor, k_mn: Optional[torch.Tensor], k_mn_zy: Optional[torch.Tensor], kmm_chol: torch.Tensor, X: Optional[torch.Tensor], t: Optional[int], trace_type: str):
    """ Nystrom kernel trace forward """
    if trace_type == 'ste':
        assert k_mn_zy is not None and t is not None, 'Incorrect arguments to trace_fwd'
        solve1 = torch.triangular_solve(k_mn_zy[:, :t], kmm_chol, upper=False, transpose=False).solution
        solve2 = torch.triangular_solve(solve1, kmm_chol, upper=False, transpose=True).solution.contiguous()
        init_val -= solve1.square_().sum(0).mean()
    elif trace_type == 'direct':
        assert k_mn is not None, 'Incorrect arguments to trace_fwd'
        solve1 = trsm(k_mn, kmm_chol, 1.0, lower=True, transpose=False)
        solve2 = trsm(solve1, kmm_chol, 1.0, lower=True, transpose=True)
        init_val -= solve1.square_().sum()
    elif trace_type == 'fast':
        assert k_mn_zy is not None and t is not None, 'Incorrect arguments to trace_fwd'
        k_subs = k_mn_zy
        assert k_subs.shape == (kmm_chol.shape[0], t), 'Shape incorrect'
        solve1 = torch.triangular_solve(k_subs, kmm_chol, upper=False, transpose=False).solution
        solve2 = torch.triangular_solve(solve1, kmm_chol, upper=False, transpose=True).solution.contiguous()
        norm = X.shape[0] / t
        init_val -= solve1.square_().sum() * norm
    else:
        raise ValueError('Trace-type %s unknown' % trace_type)
    return init_val, solve2
//  [python]
def calc_trace_fwd(init_val: torch.Tensor, k_mn: Optional[torch.Tensor], k_mn_zy: Optional[torch.Tensor], kmm_chol: torch.Tensor, X: Optional[torch.Tensor], t: Optional[int], trace_type: str):
    """ Nystrom kernel trace forward """
    if trace_type == 'ste':
        assert k_mn_zy is not None and t is not None, 'Incorrect arguments to trace_fwd'
        solve1 = torch.linalg.solve_triangular(k_mn_zy[:, :t], kmm_chol, upper=False, left=True)
        solve2 = torch.linalg.solve_triangular(solve1, kmm_chol, upper=False, left=False)
        init_val -= solve1.square_().sum(0).mean()
    elif trace_type == 'direct':
        assert k_mn is not None, 'Incorrect arguments to trace_fwd'
        solve1 = trsm(k_mn, kmm_chol, 1.0, lower=True, transpose=False)
        solve2 = trsm(solve1, kmm_chol, 1.0, lower=True, transpose=True)
        init_val -= solve1.square_().sum()
    elif trace_type == 'fast':
        assert k_mn_zy is not None and t is not None, 'Incorrect arguments to trace_fwd'
        k_subs = k_mn_zy
        assert k_subs.shape == (kmm_chol.shape[0], t), 'Shape incorrect'
        solve1 = torch.linalg.solve_triangular(k_subs, kmm_chol, upper=False, left=True)
        solve2 = torch.linalg.solve_triangular(solve1, kmm_chol, upper=False, left=False)
        norm = X.shape[0] / t
        init_val -= solve1.square_().sum() * norm
    else:
        raise ValueError('Trace-type %s unknown' % trace_type)
    return init_val, solve2
//  [python]
def backward_routine(layout, device, value_dtype, index_dtype, shapes, upper, unitriangular, transpose):
    if index_dtype == torch.int32 and layout is torch.sparse_coo:
        pytest.skip("Skipping test as sparse COO tensors with int32 indices are not supported")
    if platform == "win32" and device == torch.device("cpu"):
        pytest.skip("Skipping triangular solve CPU tests as solver not implemented for Windows OS")

    _, A_shape, B_shape, A_nnz = shapes

    As1 = rand_sparse_tri(
        A_shape,
        A_nnz,
        layout,
        upper=upper,
        strict=unitriangular,
        indices_dtype=index_dtype,
        values_dtype=value_dtype,
        device=device,
    )

    Ad2 = As1.to_dense().detach().clone()  # detach and clone to create seperate graph
    Ad3 = Ad2.detach().clone()

    Bd1 = torch.rand(*B_shape, dtype=value_dtype, device=device)
    Bd2 = Bd1.detach().clone()
    Bd3 = Bd1.detach().clone()

    As1.requires_grad_()
    Ad2.requires_grad_()
    Ad3.requires_grad_()
    Bd1.requires_grad_()
    Bd2.requires_grad_()
    Bd3.requires_grad_()

    res_ref = torch.triangular_solve(Bd2, Ad2, upper=upper, unitriangular=unitriangular, transpose=transpose).solution
    res_test = sparse_triangular_solve(As1, Bd1, upper=upper, unitriangular=unitriangular, transpose=transpose)

    # Let's add another test to make sure that the transpose argument is working as epexcted:
    if transpose:
        res_test2 = torch.linalg.solve_triangular(
            Ad3.transpose(-2, -1), Bd3, upper=not upper, unitriangular=unitriangular
        )
    else:
        res_test2 = torch.linalg.solve_triangular(Ad3, Bd3, upper=upper, unitriangular=unitriangular)

    # Generate random gradients for the backward pass
    grad_output = torch.rand_like(res_test, dtype=value_dtype, device=device)

    res_ref.backward(grad_output)
    res_test.backward(grad_output)
    res_test2.backward(grad_output)

    nz_mask = As1.grad.to_dense() != 0.0

    assert torch.allclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], atol=ATOL, rtol=RTOL)
    assert torch.allclose(As1.grad.to_dense()[nz_mask], Ad3.grad[nz_mask], atol=ATOL, rtol=RTOL)

    assert torch.allclose(Bd1.grad, Bd2.grad, atol=ATOL, rtol=RTOL)
    assert torch.allclose(Bd1.grad, Bd3.grad, atol=ATOL, rtol=RTOL)
//  [python]
def backward_routine(layout, device, value_dtype, index_dtype, shapes, upper, unitriangular, transpose):
    if index_dtype == torch.int32 and layout is torch.sparse_coo:
        pytest.skip("Skipping test as sparse COO tensors with int32 indices are not supported")
    if platform == "win32" and device == torch.device("cpu"):
        pytest.skip("Skipping triangular solve CPU tests as solver not implemented for Windows OS")

    _, A_shape, B_shape, A_nnz = shapes

    As1 = rand_sparse_tri(
        A_shape,
        A_nnz,
        layout,
        upper=upper,
        strict=unitriangular,
        indices_dtype=index_dtype,
        values_dtype=value_dtype,
        device=device,
    )

    Ad2 = As1.to_dense().detach().clone()  # detach and clone to create seperate graph
    Ad3 = Ad2.detach().clone()

    Bd1 = torch.rand(*B_shape, dtype=value_dtype, device=device)
    Bd2 = Bd1.detach().clone()
    Bd3 = Bd1.detach().clone()

    As1.requires_grad_()
    Ad2.requires_grad_()
    Ad3.requires_grad_()
    Bd1.requires_grad_()
    Bd2.requires_grad_()
    Bd3.requires_grad_()

    res_ref = torch.linalg.solve_triangular(Ad2, Bd2, upper=upper, unitriangular=unitriangular, left=True, transpose=transpose)
    res_test = sparse_triangular_solve(As1, Bd1, upper=upper, unitriangular=unitriangular, transpose=transpose)

    # Let's add another test to make sure that the transpose argument is working as epexcted:
    if transpose:
        res_test2 = torch.linalg.solve_triangular(
            Ad3.transpose(-2, -1), Bd3, upper=not upper, unitriangular=unitriangular, left=True
        )
    else:
        res_test2 = torch.linalg.solve_triangular(Ad3, Bd3, upper=upper, unitriangular=unitriangular, left=True)

    # Generate random gradients for the backward pass
    grad_output = torch.rand_like(res_test, dtype=value_dtype, device=device)

    res_ref.backward(grad_output)
    res_test.backward(grad_output)
    res_test2.backward(grad_output)

    nz_mask = As1.grad.to_dense() != 0.0

    assert torch.allclose(As1.grad.to_dense()[nz_mask], Ad2.grad[nz_mask], atol=ATOL, rtol=RTOL)
    assert torch.allclose(As1.grad.to_dense()[nz_mask], Ad3.grad[nz_mask], atol=ATOL, rtol=RTOL)

    assert torch.allclose(Bd1.grad, Bd2.grad, atol=ATOL, rtol=RTOL)
    assert torch.allclose(Bd1.grad, Bd3.grad, atol=ATOL, rtol=RTOL)
//  [python]
def precision_matrix(self):
        # We use "Woodbury matrix identity" to take advantage of low rank form::
        #     inv(W @ W.T + D) = inv(D) - inv(D) @ W @ inv(C) @ W.T @ inv(D)
        # where :math:`C` is the capacitance matrix.
        Wt_Dinv = (self._unbroadcasted_cov_factor.mT
                   / self._unbroadcasted_cov_diag.unsqueeze(-2))
        A = torch.triangular_solve(Wt_Dinv, self._capacitance_tril, upper=False)[0]
        precision_matrix = torch.diag_embed(self._unbroadcasted_cov_diag.reciprocal()) - A.mT @ A
        return precision_matrix.expand(self._batch_shape + self._event_shape +
                                       self._event_shape)
//  [python]
def precision_matrix(self):
        # We use "Woodbury matrix identity" to take advantage of low rank form::
        #     inv(W @ W.T + D) = inv(D) - inv(D) @ W @ inv(C) @ W.T @ inv(D)
        # where :math:`C` is the capacitance matrix.
        Wt_Dinv = (self._unbroadcasted_cov_factor.mT
                   / self._unbroadcasted_cov_diag.unsqueeze(-2))
        A = torch.linalg.solve_triangular(self._capacitance_tril, Wt_Dinv, upper=False)
        precision_matrix = torch.diag_embed(self._unbroadcasted_cov_diag.reciprocal()) - A.mT @ A
        return precision_matrix.expand(self._batch_shape + self._event_shape +
                                       self._event_shape)
//  [python]
def _triangular_inverse(A, upper=False):
    eye = torch.eye(A.size(-1), dtype=A.dtype, device=A.device)
    return torch.triangular_solve(eye, A, upper=upper)[0]
//  [python]
def _triangular_inverse(A, upper=False):
    eye = torch.eye(A.size(-1), dtype=A.dtype, device=A.device)
    return torch.linalg.solve_triangular(A, eye, upper=upper)
//  [python]
def backward(ctx, grad_output):
        W, L = ctx.saved_tensors
        LmT = L.mT.contiguous()
        gB = torch.triangular_solve(grad_output, LmT, upper=True)[0]
        gA = (-gB @ W.mT).tril()
        gS = (LmT @ gA).tril()
        gS = gS + gS.tril(-1).mT
        gS = torch.triangular_solve(gS, LmT, upper=True)[0]
        gX = gS @ W + gB
        return gX
//  [python]
def backward(ctx, grad_output):
        W, L = ctx.saved_tensors
        LmT = L.mT.contiguous()
        gB = torch.linalg.solve_triangular(LmT, grad_output, upper=True)
        gA = (-gB @ W.mT).tril()
        gS = (LmT @ gA).tril()
        gS = gS + gS.tril(-1).mT
        gS = torch.linalg.solve_triangular(LmT, gS, upper=True)
        gX = gS @ W + gB
        return gX
//  [python]
def backward(ctx, grad_output):
        X, W, L = ctx.saved_tensors
        gB = torch.triangular_solve(grad_output, L.mT, upper=True)[0]
        gA = (-gB @ W.mT).tril()
        gS = (L.mT @ gA).tril()
        gS = gS + gS.tril(-1).mT
        gS = torch.triangular_solve(gS, L.mT, upper=True)[0]
        gS = torch.triangular_solve(gS, L, upper=False, left=False)[0]
        gX = gS @ X + gB
        return gX
//  [python]
def backward(ctx, grad_output):
        X, W, L = ctx.saved_tensors
        gB = torch.linalg.solve_triangular(L.mT, grad_output, upper=True)
        gA = (-gB @ W.mT).tril()
        gS = (L.mT @ gA).tril()
        gS = gS + gS.tril(-1).mT
        gS = torch.linalg.solve_triangular(L.mT, gS, upper=True)
        gS = torch.linalg.solve_triangular(L, gS, upper=False, left=False)
        gX = gS @ X + gB
        return gX
//  [python]
def _kl_multivariatenormal_multivariatenormal(p, q):
    # From https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Kullback%E2%80%93Leibler_divergence
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two Multivariate Normals with\
                          different event shapes cannot be computed")

    half_term1 = (q._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1) -
                  p._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1))
    combined_batch_shape = torch._C._infer_size(q._unbroadcasted_scale_tril.shape[:-2],
                                                p._unbroadcasted_scale_tril.shape[:-2])
    n = p.event_shape[0]
    q_scale_tril = q._unbroadcasted_scale_tril.expand(combined_batch_shape + (n, n))
    p_scale_tril = p._unbroadcasted_scale_tril.expand(combined_batch_shape + (n, n))
    term2 = _batch_trace_XXT(torch.linalg.solve_triangular(q_scale_tril, p_scale_tril, upper=False))
    term3 = _batch_mahalanobis(q._unbroadcasted_scale_tril, (q.loc - p.loc))
    return half_term1 + 0.5 * (term2 + term3 - n)
//  [python]
def _kl_multivariatenormal_lowrankmultivariatenormal(p, q):
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two (Low Rank) Multivariate Normals with\
                          different event shapes cannot be computed")

    term1 = (_batch_lowrank_logdet(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                   q._capacitance_tril) -
             2 * p._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1))
    term3 = _batch_lowrank_mahalanobis(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                       q.loc - p.loc,
                                       q._capacitance_tril)
    # Expands term2 according to
    # inv(qcov) @ pcov = [inv(qD) - inv(qD) @ qW @ inv(qC) @ qW.T @ inv(qD)] @ p_tril @ p_tril.T
    #                  = [inv(qD) - A.T @ A] @ p_tril @ p_tril.T
    qWt_qDinv = (q._unbroadcasted_cov_factor.mT /
                 q._unbroadcasted_cov_diag.unsqueeze(-2))
    A = torch.triangular_solve(qWt_qDinv, q._capacitance_tril, upper=False)[0]
    term21 = _batch_trace_XXT(p._unbroadcasted_scale_tril *
                              q._unbroadcasted_cov_diag.rsqrt().unsqueeze(-1))
    term22 = _batch_trace_XXT(A.matmul(p._unbroadcasted_scale_tril))
    term2 = term21 - term22
    return 0.5 * (term1 + term2 + term3 - p.event_shape[0])
//  [python]
def _kl_multivariatenormal_lowrankmultivariatenormal(p, q):
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two (Low Rank) Multivariate Normals with\
                          different event shapes cannot be computed")

    term1 = (_batch_lowrank_logdet(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                   q._capacitance_tril) -
             2 * p._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1))
    term3 = _batch_lowrank_mahalanobis(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                       q.loc - p.loc,
                                       q._capacitance_tril)
    # Expands term2 according to
    # inv(qcov) @ pcov = [inv(qD) - inv(qD) @ qW @ inv(qC) @ qW.T @ inv(qD)] @ p_tril @ p_tril.T
    #                  = [inv(qD) - A.T @ A] @ p_tril @ p_tril.T
    qWt_qDinv = (q._unbroadcasted_cov_factor.mT /
                 q._unbroadcasted_cov_diag.unsqueeze(-2))
    A = torch.linalg.solve_triangular(q._capacitance_tril, qWt_qDinv, upper=False)
    term21 = _batch_trace_XXT(p._unbroadcasted_scale_tril *
                              q._unbroadcasted_cov_diag.rsqrt().unsqueeze(-1))
    term22 = _batch_trace_XXT(A.matmul(p._unbroadcasted_scale_tril))
    term2 = term21 - term22
    return 0.5 * (term1 + term2 + term3 - p.event_shape[0])
//  [python]
def _kl_lowrankmultivariatenormal_lowrankmultivariatenormal(p, q):
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two Low Rank Multivariate Normals with\
                          different event shapes cannot be computed")

    term1 = (_batch_lowrank_logdet(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                   q._capacitance_tril) -
             _batch_lowrank_logdet(p._unbroadcasted_cov_factor, p._unbroadcasted_cov_diag,
                                   p._capacitance_tril))
    term3 = _batch_lowrank_mahalanobis(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                       q.loc - p.loc,
                                       q._capacitance_tril)
    # Expands term2 according to
    # inv(qcov) @ pcov = [inv(qD) - inv(qD) @ qW @ inv(qC) @ qW.T @ inv(qD)] @ (pW @ pW.T + pD)
    #                  = [inv(qD) - A.T @ A] @ (pD + pW @ pW.T)
    qWt_qDinv = (q._unbroadcasted_cov_factor.mT /
                 q._unbroadcasted_cov_diag.unsqueeze(-2))
    A = torch.triangular_solve(qWt_qDinv, q._capacitance_tril, upper=False)[0]
    term21 = (p._unbroadcasted_cov_diag / q._unbroadcasted_cov_diag).sum(-1)
    term22 = _batch_trace_XXT(p._unbroadcasted_cov_factor *
                              q._unbroadcasted_cov_diag.rsqrt().unsqueeze(-1))
    term23 = _batch_trace_XXT(A * p._unbroadcasted_cov_diag.sqrt().unsqueeze(-2))
    term24 = _batch_trace_XXT(A.matmul(p._unbroadcasted_cov_factor))
    term2 = term21 + term22 - term23 - term24
    return 0.5 * (term1 + term2 + term3 - p.event_shape[0])
//  [python]
def _kl_lowrankmultivariatenormal_lowrankmultivariatenormal(p, q):
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two Low Rank Multivariate Normals with\
                          different event shapes cannot be computed")

    term1 = (_batch_lowrank_logdet(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                   q._capacitance_tril) -
             _batch_lowrank_logdet(p._unbroadcasted_cov_factor, p._unbroadcasted_cov_diag,
                                   p._capacitance_tril))
    term3 = _batch_lowrank_mahalanobis(q._unbroadcasted_cov_factor, q._unbroadcasted_cov_diag,
                                       q.loc - p.loc,
                                       q._capacitance_tril)
    # Expands term2 according to
    # inv(qcov) @ pcov = [inv(qD) - inv(qD) @ qW @ inv(qC) @ qW.T @ inv(qD)] @ (pW @ pW.T + pD)
    #                  = [inv(qD) - A.T @ A] @ (pD + pW @ pW.T)
    qWt_qDinv = (q._unbroadcasted_cov_factor.mT /
                 q._unbroadcasted_cov_diag.unsqueeze(-2))
    A = torch.linalg.solve_triangular(q._capacitance_tril, qWt_qDinv, upper=False)
    term21 = (p._unbroadcasted_cov_diag / q._unbroadcasted_cov_diag).sum(-1)
    term22 = _batch_trace_XXT(p._unbroadcasted_cov_factor *
                              q._unbroadcasted_cov_diag.rsqrt().unsqueeze(-1))
    term23 = _batch_trace_XXT(A * p._unbroadcasted_cov_diag.sqrt().unsqueeze(-2))
    term24 = _batch_trace_XXT(A.matmul(p._unbroadcasted_cov_factor))
    term2 = term21 + term22 - term23 - term24
    return 0.5 * (term1 + term2 + term3 - p.event_shape[0])
//  [python]
def _kl_lowrankmultivariatenormal_multivariatenormal(p, q):
    if p.event_shape != q.event_shape:
        raise ValueError("KL-divergence between two (Low Rank) Multivariate Normals with\
                          different event shapes cannot be computed")

    term1 = (2 * q._unbroadcasted_scale_tril.diagonal(dim1=-2, dim2=-1).log().sum(-1) -
             _batch_lowrank_logdet(p._unbroadcasted_cov_factor, p._unbroadcasted_cov_diag,
                                   p._capacitance_tril))
    term3 = _batch_mahalanobis(q._unbroadcasted_scale_tril, (q.loc - p.loc))
    # Expands term2 according to
    # inv(qcov) @ pcov = inv(q_tril @ q_tril.T) @ (pW @ pW.T + pD)
    combined_batch_shape = torch._C._infer_size(q._unbroadcasted_scale_tril.shape[:-2],
                                                p._unbroadcasted_cov_factor.shape[:-2])
    n = p.event_shape[0]
    q_scale_tril = q._unbroadcasted_scale_tril.expand(combined_batch_shape + (n, n))
    p_cov_factor = p._unbroadcasted_cov_factor.expand(combined_batch_shape +
                                                      (n, p.cov_factor.size(-1)))
    p_cov_diag = (torch.diag_embed(p._unbroadcasted_cov_diag.sqrt())
                  .expand(combined_batch_shape + (n, n)))
    term21 = _batch_trace_XXT(torch.linalg.solve_triangular(q_scale_tril, p_cov_factor, upper=False))
    term22 = _batch_trace_XXT(torch.linalg.solve_triangular(q_scale_tril, p_cov_diag, upper=False))
    term2 = term21 + term22
    return 0.5 * (term1 + term2 + term3 - p.event_shape[0])
//  [python]
def compute_samples(
        self,
        linear_solver: LinearSolver = None,
        n_samples: int = 10,
        temperature: float = 1.0,
    ) -> torch.Tensor:
        # When samples are not available, return None. This makes the outer learning loop default
        # to a perceptron loss using the mean trajectory solution from the optimizer.
        if linear_solver is None:
            return None

        # Sampling from multivariate normal using a Cholesky decomposition of AtA,
        # http://www.statsathome.com/2018/10/19/sampling-from-multivariate-normal-precision-and-covariance-parameterizations/
        delta = linear_solver.solve()
        AtA = linear_solver.linearization.hessian_approx() / temperature
        sqrt_AtA = torch.linalg.cholesky(AtA).permute(0, 2, 1)

        batch_size, n_vars = delta.shape
        y = torch.normal(
            mean=torch.zeros((n_vars, n_samples), device=delta.device),
            std=torch.ones((n_vars, n_samples), device=delta.device),
        )
        delta_samples = torch.triangular_solve(y, sqrt_AtA, upper=True).solution + (
            delta.unsqueeze(-1)
        ).repeat(1, 1, n_samples)

        x_samples = torch.zeros((batch_size, n_vars, n_samples), device=delta.device)
        for sidx in range(0, n_samples):
            var_idx = 0
            for var in linear_solver.linearization.ordering:
                new_var = var.retract(
                    delta_samples[:, var_idx : var_idx + var.dof(), sidx]
                )
                x_samples[:, var_idx : var_idx + var.dof(), sidx] = new_var.tensor
                var_idx = var_idx + var.dof()

        return x_samples
//  [python]
def compute_samples(
        self,
        linear_solver: LinearSolver = None,
        n_samples: int = 10,
        temperature: float = 1.0,
    ) -> torch.Tensor:
        # When samples are not available, return None. This makes the outer learning loop default
        # to a perceptron loss using the mean trajectory solution from the optimizer.
        if linear_solver is None:
            return None

        # Sampling from multivariate normal using a Cholesky decomposition of AtA,
        # http://www.statsathome.com/2018/10/19/sampling-from-multivariate-normal-precision-and-covariance-parameterizations/
        delta = linear_solver.solve()
        AtA = linear_solver.linearization.hessian_approx() / temperature
        sqrt_AtA = torch.linalg.cholesky(AtA).permute(0, 2, 1)

        batch_size, n_vars = delta.shape
        y = torch.normal(
            mean=torch.zeros((n_vars, n_samples), device=delta.device),
            std=torch.ones((n_vars, n_samples), device=delta.device),
        )
        delta_samples = (torch.linalg.solve_triangular(sqrt_AtA, y, upper=True)) + (
            delta.unsqueeze(-1)
        ).repeat(1, 1, n_samples)

        x_samples = torch.zeros((batch_size, n_vars, n_samples), device=delta.device)
        for sidx in range(0, n_samples):
            var_idx = 0
            for var in linear_solver.linearization.ordering:
                new_var = var.retract(
                    delta_samples[:, var_idx : var_idx + var.dof(), sidx]
                )
                x_samples[:, var_idx : var_idx + var.dof(), sidx] = new_var.tensor
                var_idx = var_idx + var.dof()

        return x_samples
//  [python]
def calculate_W_prefactor(S, P, L, U):
    '''
    Our new Biharmonic weight formulation with prefactorization
    '''
    # make complementary selector matrix T^T (n x n) w.r.t current S
    pred_idx = torch.argmax(S, dim=-1)
    index = torch.ones((1, U.shape[-1]), dtype=torch.float64).to(S.device)
    index = index.scatter_(index=pred_idx, dim=-1, value=0.0)
    # T^T T
    TT = torch.diag_embed(index).to(torch.float64)

    b = torch.matmul(P, torch.transpose(S, -2, -1))
    y = torch.triangular_solve(b, L, upper=False).solution
    Z = torch.triangular_solve(y, U, upper=True).solution

    X = torch.solve(torch.bmm(S,Z).transpose(-2,-1), Z.transpose(-2,-1))[0].transpose(-2,-1)
    W_reformulated_full = torch.transpose(S, -2, -1).to(torch.float64) + torch.matmul(TT, X).to(torch.float64)

    return W_reformulated_full
//  [python]
def calculate_W_prefactor(S, P, L, U):
    '''
    Our new Biharmonic weight formulation with prefactorization
    '''
    # make complementary selector matrix T^T (n x n) w.r.t current S
    pred_idx = torch.argmax(S, dim=-1)
    index = torch.ones((1, U.shape[-1]), dtype=torch.float64).to(S.device)
    index = index.scatter_(index=pred_idx, dim=-1, value=0.0)
    # T^T T
    TT = torch.diag_embed(index).to(torch.float64)

    b = torch.matmul(P, torch.transpose(S, -2, -1))
    y = torch.linalg.solve_triangular(L, b, upper=False)
    Z = torch.linalg.solve_triangular(U, y, upper=True)

    X = torch.linalg.solve(torch.bmm(S,Z).transpose(-2,-1), Z.transpose(-2,-1)).transpose(-2,-1)
    W_reformulated_full = torch.transpose(S, -2, -1).to(torch.float64) + torch.matmul(TT, X).to(torch.float64)

    return W_reformulated_full
//  [python]
def solve_triangular(
    x1: torch.Tensor,
    x2: torch.Tensor,
    /,
    *,
    upper: bool = True,
    adjoint: bool = False,
    unit_diagonal: bool = False,
    out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    if adjoint:
        x1 = torch.adjoint(x1)
        upper = not upper
    return torch.triangular_solve(x2, x1, upper=upper, unitriangular=unit_diagonal, out=out)[0]
//  [python]
def solve_triangular(
    x1: torch.Tensor,
    x2: torch.Tensor,
    /,
    *,
    upper: bool = True,
    adjoint: bool = False,
    unit_diagonal: bool = False,
    out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    if adjoint:
        x1 = torch.adjoint(x1)
        upper = not upper
    return torch.linalg.solve_triangular(
        x1, x2, upper=upper, unitriangular=unit_diagonal, out=out
    )
//  [python]
def weight_inverse(self):
        """Cost:
            inverse = O(D^3)
        where:
            D = num of features
        """
        lower, upper = self._create_lower_upper()
        identity = torch.eye(
            self.features, self.features, device=self.lower_entries.device)
        lower_inverse = torch.triangular_solve(
            identity, lower, upper=False, unitriangular=True
        )[0]
        weight_inverse = torch.triangular_solve(
            lower_inverse, upper, upper=True, unitriangular=False
        )[0]
        return weight_inverse
//  [python]
def weight_inverse(self):
        """Cost:
            inverse = O(D^3)
        where:
            D = num of features
        """
        lower, upper = self._create_lower_upper()
        identity = torch.eye(
            self.features, self.features, device=self.lower_entries.device)
        lower_inverse = torch.linalg.solve_triangular(
            lower, identity, upper=False, unitriangular=True
        )
        weight_inverse = torch.linalg.solve_triangular(
            upper, lower_inverse, upper=True, unitriangular=False
        )
        return weight_inverse
//  [python]
def solve_triangular_torch(A, b, lower=False, trans=0, **kwargs):
    """Solve a linear system using Pytorch."""
    import torch

    A = torch.transpose(torch.tensor(A), -2, -1) if trans == 1 or trans == "T" else torch.tensor(A)
    return torch.triangular_solve(
        b if isinstance(b, torch.Tensor) else torch.tensor(b), A, upper=not lower, **kwargs
    )[0]
//  [python]
def solve_triangular_torch(A, b, lower=False, trans=0, **kwargs):
    """Solve a linear system using Pytorch."""
    import torch

    A = torch.transpose(torch.tensor(A), -2, -1) if trans == 1 or trans == "T" else torch.tensor(A)
    return torch.linalg.solve_triangular(
        A, b if isinstance(b, torch.Tensor) else torch.tensor(b), upper=not lower, **kwargs
    )
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^2N)
            logabsdet = O(D)
        where:
            D = num of features
            N = num of inputs
        """
        lower, upper = self._create_lower_upper()
        outputs = inputs - self.bias
        outputs = torch.triangular_solve(
            outputs.t(), lower, upper=False, unitriangular=True
        )[0]
        outputs = torch.triangular_solve(
            outputs, upper, upper=True, unitriangular=False
        )[0]
        outputs = outputs.t()

        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * inputs.new_ones(outputs.shape[0])

        return outputs, logabsdet
//  [python]
def inverse_no_cache(self, inputs):
        """Cost:
            output = O(D^2N)
            logabsdet = O(D)
        where:
            D = num of features
            N = num of inputs
        """
        lower, upper = self._create_lower_upper()
        outputs = inputs - self.bias
        outputs = torch.linalg.solve_triangular(
            lower, outputs.t(), upper=False, unitriangular=True
        )
        outputs = torch.linalg.solve_triangular(
            upper, outputs, upper=True, unitriangular=False
        )
        outputs = outputs.t()

        logabsdet = -self.logabsdet()
        logabsdet = logabsdet * inputs.new_ones(outputs.shape[0])

        return outputs, logabsdet
//  [python]
def linearize_sparse_depth_prior(L_mm):
  B, N_train, _ = L_mm.shape
  device = L_mm.device

  da_ds = -torch.ones(B, N_train, 1, device=device)
  dr_ds = torch.triangular_solve(da_ds, L_mm, upper=False).solution.unsqueeze(2)
  da_dd = torch.eye(N_train, device=device).reshape((1,N_train,N_train)).repeat(B,1,1)
  dr_dd = torch.triangular_solve(da_dd, L_mm, upper=False).solution.unsqueeze(2)

  H_s_s = get_hessian_diag_block(dr_ds)
  H_s_d = get_hessian_off_diag_block(dr_ds, dr_dd)
  H_d_d = get_hessian_diag_block(dr_dd)

  return dr_ds, dr_dd, H_s_s, H_s_d, H_d_d
//  [python]
def linearize_sparse_depth_prior(L_mm):
  B, N_train, _ = L_mm.shape
  device = L_mm.device

  da_ds = -torch.ones(B, N_train, 1, device=device)
  dr_ds = torch.linalg.solve_triangular(L_mm, da_ds, upper=False).unsqueeze(2)
  da_dd = torch.eye(N_train, device=device).reshape((1,N_train,N_train)).repeat(B,1,1)
  dr_dd = torch.linalg.solve_triangular(L_mm, da_dd, upper=False).unsqueeze(2)

  H_s_s = get_hessian_diag_block(dr_ds)
  H_s_d = get_hessian_off_diag_block(dr_ds, dr_dd)
  H_d_d = get_hessian_diag_block(dr_dd)

  return dr_ds, dr_dd, H_s_s, H_s_d, H_d_d
//  [python]
def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.unsqueeze(2).unsqueeze(3)
        m = x.mean(0).view(self.num_features, -1).mean(-1).view(1, -1, 1, 1)
        if not self.training and self.track_running_stats:  # for inference
            m = self.running_mean
        xn = x - m

        T = xn.permute(1, 0, 2, 3).contiguous().view(self.num_features, -1)
        f_cov = torch.mm(T, T.permute(1, 0)) / (T.shape[-1] - 1)

        eye = torch.eye(self.num_features).type(f_cov.type())

        if not self.training and self.track_running_stats:  # for inference
            f_cov = self.running_variance

        f_cov_shrinked = (1 - self.eps) * f_cov + self.eps * eye

        inv_sqrt = torch.solve(eye, torch.cholesky(f_cov_shrinked))[0]

        inv_sqrt = inv_sqrt.contiguous().view(
            self.num_features, self.num_features, 1, 1
        )

        decorrelated = F.conv2d(xn, inv_sqrt)

        if self.training and self.track_running_stats:
            self.running_mean = torch.add(
                self.momentum * m.detach(),
                (1 - self.momentum) * self.running_mean,
                out=self.running_mean,
            )
            self.running_variance = torch.add(
                self.momentum * f_cov.detach(),
                (1 - self.momentum) * self.running_variance,
                out=self.running_variance,
            )

        return decorrelated.squeeze(2).squeeze(2)
//  [python]
def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x.unsqueeze(2).unsqueeze(3)
        m = x.mean(0).view(self.num_features, -1).mean(-1).view(1, -1, 1, 1)
        if not self.training and self.track_running_stats:  # for inference
            m = self.running_mean
        xn = x - m

        T = xn.permute(1, 0, 2, 3).contiguous().view(self.num_features, -1)
        f_cov = torch.mm(T, T.permute(1, 0)) / (T.shape[-1] - 1)

        eye = torch.eye(self.num_features).type(f_cov.type())

        if not self.training and self.track_running_stats:  # for inference
            f_cov = self.running_variance

        f_cov_shrinked = (1 - self.eps) * f_cov + self.eps * eye

        inv_sqrt = torch.linalg.solve_triangular(
            torch.linalg.cholesky(f_cov_shrinked), eye, upper=False
        )

        inv_sqrt = inv_sqrt.contiguous().view(
            self.num_features, self.num_features, 1, 1
        )

        decorrelated = F.conv2d(xn, inv_sqrt)

        if self.training and self.track_running_stats:
            self.running_mean = torch.add(
                self.momentum * m.detach(),
                (1 - self.momentum) * self.running_mean,
                out=self.running_mean,
            )
            self.running_variance = torch.add(
                self.momentum * f_cov.detach(),
                (1 - self.momentum) * self.running_variance,
                out=self.running_variance,
            )

        return decorrelated.squeeze(2).squeeze(2)
//  [python]
def _get_score(self, model, x, i, L=None, **kwargs):
        """
        Returns score: ∇_xs log p(xs)
        """
        if isinstance(i, float):
            i = torch.Tensor([i]).to(x)
        if i.shape[:-1] != x.shape[:-1]:
            i = i.view(*(1,) * len(x.shape)).expand_as(x[...,:1])

        beta_int = self.beta_fn.integral(i)
        std = (1 - torch.exp(-beta_int)).clamp(1e-5).sqrt()

        noise = model(x, i=i, **kwargs)

        if L is not None:
            # We have to compute the score using -Sigma.inv() @ noise / std
            # assuming noise~N(0, Sigma).
            # If `predict_gaussian_noise=False`, compute (LL^T).inv()
            # Else, we can simplify (LL^T).inv() @ L @ noise
            # to (L^T).inv() @ noise, where noise~N(0, I).
            # So we anyways have to do (L^T).inv(), and sometimes L.inv()
            if not self.predict_gaussian_noise:
                noise = torch.triangular_solve(noise, L, upper=False)[0]
            noise = torch.triangular_solve(noise, L.transpose(-1, -2), upper=True)[0]

        score = -noise / std
        return score