//  [python]
def forward(self) -> torch.Tensor:
        # get sigma/mean or each level
        sigma_w, mu_w = self.implied_sigma_mu()
        sigma_l2, mu_l2 = self.implied_sigma_mu(suffix="_l2")

        # decompose into 11, 12, 21, 22
        sigma_b, sigma_xx, sigma_yx, mu_b, mu_x = self._split(sigma_l2, mu_l2)

        # cluster FIML -2 * logL (without constants)
        loss = torch.zeros(1, dtype=mu_w.dtype, device=mu_w.device)

        # go through each cluster separately
        data_ys_available = ~torch.isnan(self.data_ys)
        cache_S_ij = {}
        cache_S_j_R_j = {}
        sigma_b_logdet = None
        sigma_b_inv = None
        if not self.naive_implementation:
            L_sigma_b = torch.cholesky(sigma_b)
            sigma_b_logdet = L_sigma_b.diagonal().log().sum() * 2
            sigma_b_inv = torch.cholesky_inverse(L_sigma_b)
        for cluster_slice, batches in self.missing_patterns:
            # get cluster data and define R_j for current cluster j
            cluster_x = self.data_xs[cluster_slice.start, :]
            R_j_index = ~torch.isnan(cluster_x)
            no_cluster = ~R_j_index.any()

            # cache
            key = (
                tuple(R_j_index.tolist()),
                tuple(tuple(x) for x in data_ys_available[cluster_slice, :].tolist()),
            )
            sigma_j_logdet, sigma_j_inv = cache_S_j_R_j.get(key, (None, None))

            # define S_ij and S_j
            S_ijs = []
            eye_w = torch.eye(mu_w.shape[0], dtype=mu_w.dtype, device=mu_w.device)
            lambda_ijs_logdet_sum = 0.0
            lambda_ijs_inv = []
            A_j = torch.zeros_like(sigma_w)
            for batch_slice in batches:
                size = batch_slice.stop - batch_slice.start
                available = data_ys_available[batch_slice.start]
                S_ij = eye_w[available, :]
                S_ijs.extend([S_ij] * size)

                if self.naive_implementation or sigma_j_logdet is not None:
                    continue

                key_S_ij = tuple(available.tolist())
                lambda_ij_inv, lambda_ij_logdet, a_j = cache_S_ij.get(
                    key_S_ij, (None, None, None),
                )

                if lambda_ij_inv is None:
                    lambda_ij = sigma_w  # no missing data
                    if S_ij.shape[0] != eye_w.shape[0]:
                        # missing data
                        lambda_ij = S_ij.mm(sigma_w.mm(S_ij.t()))

                    L_lambda_ij = torch.cholesky(lambda_ij)
                    lambda_ij_logdet = L_lambda_ij.diagonal().log().sum() * 2
                    lambda_ij_inv = torch.cholesky_inverse(L_lambda_ij)

                    if S_ij.shape[0] != eye_w.shape[0]:
                        # missing data
                        a_j = S_ij.t().mm(torch.cholesky_solve(S_ij, L_lambda_ij))
                    else:
                        a_j = lambda_ij_inv
                    cache_S_ij[key_S_ij] = lambda_ij_inv, lambda_ij_logdet, a_j

                lambda_ijs_inv.extend([lambda_ij_inv] * size)
                lambda_ijs_logdet_sum = lambda_ijs_logdet_sum + lambda_ij_logdet * size
                A_j = A_j + a_j * size

            S_j = torch.cat(S_ijs, dim=0)

            # means
            y_j = torch.cat(
                [
                    self.data_ys[cluster_slice, :][data_ys_available[cluster_slice, :]][
                        :, None,
                    ],
                    cluster_x[R_j_index, None],
                ],
            )
            mu_y = mu_w + mu_b
            mu_j = torch.cat([S_j.mm(mu_y), mu_x[R_j_index]])
            mean_diff = y_j - mu_j
            G_yj = mean_diff.mm(mean_diff.t())

            if sigma_j_logdet is None and not self.naive_implementation:
                sigma_b_inv_A_j = sigma_b_inv + A_j
                B_j = torch.inverse(sigma_b_inv_A_j)
                C_j = eye_w - A_j.mm(B_j)
                D_j = C_j.mm(A_j)
                lambda_inv = torch.block_diag(*lambda_ijs_inv)
                V_j_inv = lambda_inv - lambda_inv.mm(
                    S_j.mm(B_j.mm(S_j.t().mm(lambda_inv))),
                )

                if no_cluster:
                    # no cluster
                    sigma_11_j = V_j_inv
                    sigma_21_j = torch.empty(
                        0, device=sigma_11_j.device, dtype=sigma_11_j.dtype,
                    )
                    sigma_22_1 = torch.empty(
                        [0, 0], device=sigma_11_j.device, dtype=sigma_11_j.dtype,
                    )
                    sigma_22_inv = sigma_21_j

                else:
                    # normal case
                    sigma_22_1 = (sigma_xx - sigma_yx.t().mm(D_j.mm(sigma_yx)))[
                        R_j_index, :,
                    ][:, R_j_index]
                    sigma_22_inv = torch.inverse(sigma_22_1)
                    sigma_jyx = S_j.mm(sigma_yx[:, R_j_index])
                    sigma_11_j = (
                        V_j_inv.mm(
                            sigma_jyx.mm(sigma_22_inv.mm(sigma_jyx.t().mm(V_j_inv))),
                        )
                        + V_j_inv
                    )
                    sigma_21_j = -sigma_22_inv.mm(sigma_jyx.t().mm(V_j_inv))

                sigma_j_inv = torch.cat(
                    [
                        torch.cat([sigma_11_j, sigma_21_j]),
                        torch.cat([sigma_21_j.t(), sigma_22_inv]),
                    ],
                    dim=1,
                )

                sigma_j_logdet = (
                    lambda_ijs_logdet_sum
                    + sigma_b_logdet
                    + torch.logdet(sigma_b_inv_A_j)
                    + torch.logdet(sigma_22_1)
                )
                cache_S_j_R_j[key] = (sigma_j_logdet, sigma_j_inv)

            elif sigma_j_logdet is None:
                # naive
                sigma_j = S_j.mm(sigma_b.mm(S_j.t())) + torch.block_diag(
                    *[S_ij.mm(sigma_w.mm(S_ij.t())) for S_ij in S_ijs],
                )
                if not no_cluster:
                    sigma_j_12 = S_j.mm(sigma_yx[:, R_j_index])
                    sigma_j_21 = sigma_j_12.t()
                    sigma_j_22 = sigma_xx[R_j_index, :][:, R_j_index]
                    sigma_j = torch.cat(
                        [
                            torch.cat([sigma_j, sigma_j_21]),
                            torch.cat([sigma_j_12, sigma_j_22]),
                        ],
                        dim=1,
                    )
                L_sigma_j = torch.cholesky(sigma_j)
                sigma_j_logdet = L_sigma_j.diagonal().log().sum() * 2
                sigma_j_inv = torch.cholesky_inverse(L_sigma_j)

                cache_S_j_R_j[key] = (sigma_j_logdet, sigma_j_inv)

            loss_current = sigma_j_logdet + torch.trace(sigma_j_inv.mm(G_yj))
            loss = loss + loss_current.clamp(min=0.0)

        return loss