jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def evaluate_t(self, y_loss: str, test_data: TestDataSetTorch):
        target = test_data.structural
        with torch.no_grad():
            pred = self.predict_t(test_data.treatment, test_data.covariate)
        if y_loss == 'bin':
            return (torch.norm((target - pred)) ** 2) / target.size()[0]
        else:
            return (torch.norm((target - pred)) ** 2) / target.size()[0]
//  [python]
def evaluate_t(self, y_loss: str, test_data: TestDataSetTorch):
        target = test_data.structural
        with torch.no_grad():
            pred = self.predict_t(test_data.treatment, test_data.covariate)
        if y_loss == 'bin':
            return (torch.linalg.norm((target - pred)) ** 2) / target.size()[0]
        else:
            return (torch.linalg.norm((target - pred)) ** 2) / target.size()[0]
//  [python]
def prime(x, proj_matrix):
    # x: shape (B, len, dim)
    # proj_matrix (dim, proj_dim)
    _, m = proj_matrix.shape
    
    # Compute offset in logspace
    norm_x_squared = torch.norm(x, dim=-1).pow(2) * 0.5
    sqrt_m = 0.5 * np.log(m)
    sqrt_2 = 0.5 * np.log(2)
    offset = norm_x_squared + sqrt_m + sqrt_2
    offset = offset.unsqueeze(-1)
    check(offset, [x.shape[0], x.shape[1], 1])
    
    u = torch.matmul(x, proj_matrix)
    pos = torch.exp(u - offset)
    neg = torch.exp(-u - offset)

    # last dim is feat.
    out = torch.cat([pos, neg], dim=-1)
    return out
//  [python]
def prime(x, proj_matrix):
    # x: shape (B, len, dim)
    # proj_matrix (dim, proj_dim)
    _, m = proj_matrix.shape
    
    # Compute offset in logspace
    norm_x_squared = torch.linalg.norm(x, dim=-1).pow(2) * 0.5
    sqrt_m = 0.5 * np.log(m)
    sqrt_2 = 0.5 * np.log(2)
    offset = norm_x_squared + sqrt_m + sqrt_2
    offset = offset.unsqueeze(-1)
    check(offset, [x.shape[0], x.shape[1], 1])
    
    u = torch.matmul(x, proj_matrix)
    pos = torch.exp(u - offset)
    neg = torch.exp(-u - offset)

    # last dim is feat.
    out = torch.cat([pos, neg], dim=-1)
    return out
//  [python]
def _is_normalized(mat, dim=-1):
    """
    Check if one dim of a matrix is normalized.
    """
    norm = torch.norm(mat, p=2, dim=dim)
    return (norm - 1.0).abs().max() < EPS
//  [python]
def _is_normalized(mat, dim=-1):
    """
    Check if one dim of a matrix is normalized.
    """
    norm = torch.linalg.norm(mat, ord=2, dim=dim)
    return (norm - 1.0).abs().max() < EPS
//  [python]
def _process_zero_quat(self):
        """Convert zero-norm quat to (1, 0, 0, 0)."""
        with torch.no_grad():
            norms = torch.norm(self._rot, p=2, dim=-1, keepdim=True)
            new_rot = torch.zeros_like(self._rot)
            new_rot[..., 0] = 1.0  # zero quat
            valid_mask = (norms.abs() > 0.5).repeat_interleave(4, dim=-1)
        self._rot = torch.where(valid_mask, self._rot, new_rot)
//  [python]
def _process_zero_quat(self):
        """Convert zero-norm quat to (1, 0, 0, 0)."""
        with torch.no_grad():
            norms = torch.linalg.norm(self._rot, ord=2, dim=-1, keepdim=True)
            new_rot = torch.zeros_like(self._rot)
            new_rot[..., 0] = 1.0  # zero quat
            valid_mask = (norms.abs() > 0.5).repeat_interleave(4, dim=-1)
        self._rot = torch.where(valid_mask, self._rot, new_rot)
//  [python]
def bond_length_l1(
    pred,
    target,
    mask,
    offsets=[1, 2],
):
    losses = []
    for c in range(pred.shape[0]):
        m, p, t = mask[c], pred[c], target[c]
        for o in offsets:
            m_ = (torch.stack([m[:-o], m[o:]])).all(0)
            pred_lens = torch.norm(p[:-o] - p[o:], dim=-1)
            target_lens = torch.norm(t[:-o] - t[o:], dim=-1)

            losses.append(
                torch.abs(
                    pred_lens[m_] - target_lens[m_],
                ).mean()
                / o
            )

    return torch.stack(losses)
//  [python]
def bond_length_l1(
    pred,
    target,
    mask,
    offsets=[1, 2],
):
    losses = []
    for c in range(pred.shape[0]):
        m, p, t = mask[c], pred[c], target[c]
        for o in offsets:
            m_ = (torch.stack([m[:-o], m[o:]])).all(0)
            pred_lens = torch.linalg.norm(p[:-o] - p[o:], dim=-1)
            target_lens = torch.linalg.norm(t[:-o] - t[o:], dim=-1)

            losses.append(
                torch.abs(
                    pred_lens[m_] - target_lens[m_],
                ).mean()
                / o
            )

    return torch.stack(losses)
//  [python]
def fit_2sls(treatment_1st_feature: torch.Tensor,
                 instrumental_1st_feature: torch.Tensor,
                 instrumental_2nd_feature: torch.Tensor,
                 covariate_2nd_feature: Optional[torch.Tensor],
                 outcome_2nd_t: torch.Tensor,
                 lam1: float, lam2: float,
                 add_stage1_intercept: bool,
                 add_stage2_intercept: bool,
                 ):

        # stage1
        feature = DFIVModel.augment_stage1_feature(instrumental_1st_feature, add_stage1_intercept)
        stage1_weight = fit_linear(treatment_1st_feature, feature, lam1)

        # predicting for stage 2
        feature = DFIVModel.augment_stage1_feature(instrumental_2nd_feature,
                                                   add_stage1_intercept)
        predicted_treatment_feature = linear_reg_pred(feature, stage1_weight)

        # stage2
        feature = DFIVModel.augment_stage2_feature(predicted_treatment_feature,
                                                   covariate_2nd_feature,
                                                   add_stage2_intercept)

        stage2_weight = fit_linear(outcome_2nd_t, feature, lam2)
        pred = linear_reg_pred(feature, stage2_weight)
        stage2_loss = torch.norm((outcome_2nd_t - pred)) ** 2 + lam2 * torch.norm(stage2_weight) ** 2

        labels = logit(outcome_2nd_t)
        logits = logit(pred)
        stage2_log_loss = (-(torch.log(logits) * labels+torch.log(1-logits) * (1-labels))).sum() + lam2 * torch.norm(stage2_weight) ** 2

        return dict(stage1_weight=stage1_weight,
                    predicted_treatment_feature=predicted_treatment_feature,
                    stage2_weight=stage2_weight,
                    stage2_loss=stage2_loss,
                    stage2_log_loss=stage2_log_loss)
//  [python]
def fit_2sls(treatment_1st_feature: torch.Tensor,
                 instrumental_1st_feature: torch.Tensor,
                 instrumental_2nd_feature: torch.Tensor,
                 covariate_2nd_feature: Optional[torch.Tensor],
                 outcome_2nd_t: torch.Tensor,
                 lam1: float, lam2: float,
                 add_stage1_intercept: bool,
                 add_stage2_intercept: bool,
                 ):

        # stage1
        feature = DFIVModel.augment_stage1_feature(instrumental_1st_feature, add_stage1_intercept)
        stage1_weight = fit_linear(treatment_1st_feature, feature, lam1)

        # predicting for stage 2
        feature = DFIVModel.augment_stage1_feature(instrumental_2nd_feature,
                                                   add_stage1_intercept)
        predicted_treatment_feature = linear_reg_pred(feature, stage1_weight)

        # stage2
        feature = DFIVModel.augment_stage2_feature(predicted_treatment_feature,
                                                   covariate_2nd_feature,
                                                   add_stage2_intercept)

        stage2_weight = fit_linear(outcome_2nd_t, feature, lam2)
        pred = linear_reg_pred(feature, stage2_weight)
        stage2_loss = torch.linalg.norm((outcome_2nd_t - pred)) ** 2 + lam2 * torch.linalg.norm(stage2_weight) ** 2

        labels = logit(outcome_2nd_t)
        logits = logit(pred)
        stage2_log_loss = (-(torch.log(logits) * labels+torch.log(1-logits) * (1-labels))).sum() + lam2 * torch.linalg.norm(stage2_weight) ** 2

        return dict(stage1_weight=stage1_weight,
                    predicted_treatment_feature=predicted_treatment_feature,
                    stage2_weight=stage2_weight,
                    stage2_loss=stage2_loss,
                    stage2_log_loss=stage2_log_loss)
//  [python]
def _calc_metric(self, matrix: torch.Tensor):
        return torch.norm(matrix, p=self.norm).item(), 1
//  [python]
def _calc_metric(self, matrix: torch.Tensor):
        return torch.linalg.norm(matrix, ord=self.norm).item(), 1
//  [python]
def subatt_score(self, embeddings):
     
      L = torch.tensor([]).to(self.device)
      for i in range(self.Gnum_m):     
          Ls = 0
          for j in range(self.Gnum_m):
            if i!=j:
              distance = torch.norm(embeddings[i]- embeddings[j], p=2, keepdim=True) ** 2 
              Ls += torch.log((distance + 1) / (distance + 1e-4))
          L = torch.cat((L, torch.sum(Ls).unsqueeze(0)), 0)  
      att_score = L / L.sum()

      return att_score.reshape(1,-1)
//  [python]
def subatt_score(self, embeddings):
     
      L = torch.tensor([]).to(self.device)
      for i in range(self.Gnum_m):     
          Ls = 0
          for j in range(self.Gnum_m):
            if i!=j:
              distance = torch.linalg.norm(embeddings[i]- embeddings[j], ord=2, keepdim=True) ** 2 
              Ls += torch.log((distance + 1) / (distance + 1e-4))
          L = torch.cat((L, torch.sum(Ls).unsqueeze(0)), 0)  
      att_score = L / L.sum()

      return att_score.reshape(1,-1)
//  [python]
def _update_required(
        self,
        positions: torch.tensor,
        cells: torch.tensor,
        idx_m: torch.tensor,
        n_molecules: int,
    ):
        """
        Use displacement and cell changes to determine, whether an update of the neighbor list is necessary.

        Args:
            positions (torch.Tensor): Atom positions.
            cells (torch.Tensor): Simulation cells.
            idx_m (torch.Tensor): Molecular indices.
            n_molecules (int): Number of molecules in simulation

        Returns:
            bool: Udate is required.
        """

        if self.previous_positions is None:
            # Everything needs to be updated
            update_required = torch.ones(n_molecules, device=idx_m.device).bool()
        elif n_molecules != len(self.molecular_indices):
            self.molecular_indices = None
            update_required = torch.ones(n_molecules, device=idx_m.device).bool()
        else:
            # Check for changes is positions
            update_positions = (
                torch.norm(self.previous_positions - positions, dim=1)
                > 0.5 * self.cutoff_shell
            ).float()

            # Map to individual molecules
            update_required = torch.zeros(n_molecules, device=idx_m.device).float()
            update_required = update_required.index_add(
                0, idx_m, update_positions
            ).bool()

            # Check for cell changes (is no cells are required, this will always be zero)
            update_cells = torch.any((self.previous_cells != cells).view(-1, 9), dim=1)
            update_required = torch.logical_or(update_required, update_cells)

        return update_required
//  [python]
def _update_required(
        self,
        positions: torch.tensor,
        cells: torch.tensor,
        idx_m: torch.tensor,
        n_molecules: int,
    ):
        """
        Use displacement and cell changes to determine, whether an update of the neighbor list is necessary.

        Args:
            positions (torch.Tensor): Atom positions.
            cells (torch.Tensor): Simulation cells.
            idx_m (torch.Tensor): Molecular indices.
            n_molecules (int): Number of molecules in simulation

        Returns:
            bool: Udate is required.
        """

        if self.previous_positions is None:
            # Everything needs to be updated
            update_required = torch.ones(n_molecules, device=idx_m.device).bool()
        elif n_molecules != len(self.molecular_indices):
            self.molecular_indices = None
            update_required = torch.ones(n_molecules, device=idx_m.device).bool()
        else:
            # Check for changes is positions
            update_positions = (
                torch.linalg.norm(self.previous_positions - positions, dim=1)
                > 0.5 * self.cutoff_shell
            ).float()

            # Map to individual molecules
            update_required = torch.zeros(n_molecules, device=idx_m.device).float()
            update_required = update_required.index_add(
                0, idx_m, update_positions
            ).bool()

            # Check for cell changes (is no cells are required, this will always be zero)
            update_cells = torch.any((self.previous_cells != cells).view(-1, 9), dim=1)
            update_required = torch.logical_or(update_required, update_cells)

        return update_required
//  [python]
def compute_reconstruction_error(tensor, target_tensor, normalize=False):
    reconstruction_error = torch.norm(tensor - target_tensor, p="fro").item()
    if normalize:
        target_norm = torch.norm(target_tensor, p="fro").item()
        if target_norm != 0:
            reconstruction_error /= target_norm

    return reconstruction_error
//  [python]
def compute_reconstruction_error(tensor, target_tensor, normalize=False):
    reconstruction_error = torch.linalg.norm(tensor - target_tensor, ord="fro").item()
    if normalize:
        target_norm = torch.linalg.norm(target_tensor, ord="fro").item()
        if target_norm != 0:
            reconstruction_error /= target_norm

    return reconstruction_error
//  [python]
def exp_SO3(phi):
    phi_norm = torch.norm(phi)

    if phi_norm > 1e-8:
        unit_phi = phi / phi_norm
        unit_phi_skewed = skew3(unit_phi)
        C = torch.eye(3, 3, device=phi.device) + torch.sin(phi_norm) * unit_phi_skewed + \
            (1 - torch.cos(phi_norm)) * torch.mm(unit_phi_skewed, unit_phi_skewed)
    else:
        phi_skewed = skew3(phi)
        C = torch.eye(3, 3, device=phi.device) + phi_skewed + 0.5 * torch.mm(phi_skewed, phi_skewed)

    return C
//  [python]
def exp_SO3(phi):
    phi_norm = torch.linalg.norm(phi)

    if phi_norm > 1e-8:
        unit_phi = phi / phi_norm
        unit_phi_skewed = skew3(unit_phi)
        C = torch.eye(3, 3, device=phi.device) + torch.sin(phi_norm) * unit_phi_skewed + \
            (1 - torch.cos(phi_norm)) * torch.mm(unit_phi_skewed, unit_phi_skewed)
    else:
        phi_skewed = skew3(phi)
        C = torch.eye(3, 3, device=phi.device) + phi_skewed + 0.5 * torch.mm(phi_skewed, phi_skewed)

    return C
//  [python]
def embed(self, entity_mentions: List[str]) -> Dict[str, np.ndarray]:
        query_embeddings: Dict[str, List] = {"dense": []}

        inputs = [Sentence(name) for name in entity_mentions]

        with torch.no_grad():
            for start in range(0, len(entity_mentions), self.batch_size):
                end = min(start + self.batch_size, len(entity_mentions))
                batch = inputs[start:end]
                self.embeddings["dense"].embed(batch)
                for sent in batch:
                    emb = sent.get_embedding(self.embeddings["dense"].get_names())
                    if self.similarity_metric == SimilarityMetric.COSINE:
                        emb = emb / torch.norm(emb)
                    query_embeddings["dense"].append(emb.cpu().numpy())
                    sent.clear_embeddings(self.embeddings["dense"].get_names())

                # Sanity conversion: if flair.device was set as a string, convert to torch.device
                if isinstance(flair.device, str):
                    flair.device = torch.device(flair.device)

                if flair.device.type == "cuda":
                    torch.cuda.empty_cache()

        if self.hybrid_search:
            query_embeddings["sparse"] = []
            self.embeddings["sparse"].embed(inputs)
            for sent in inputs:
                sparse_emb = sent.get_embedding(self.embeddings["sparse"].get_names())
                if self.similarity_metric == SimilarityMetric.COSINE:
                    sparse_emb = sparse_emb / torch.norm(sparse_emb)
                query_embeddings["sparse"].append(sparse_emb.cpu().numpy())
                sent.clear_embeddings(self.embeddings["sparse"].get_names())

        return {k: np.stack(v, axis=0) for k, v in query_embeddings.items()}
//  [python]
def embed(self, entity_mentions: List[str]) -> Dict[str, np.ndarray]:
        query_embeddings: Dict[str, List] = {"dense": []}

        inputs = [Sentence(name) for name in entity_mentions]

        with torch.no_grad():
            for start in range(0, len(entity_mentions), self.batch_size):
                end = min(start + self.batch_size, len(entity_mentions))
                batch = inputs[start:end]
                self.embeddings["dense"].embed(batch)
                for sent in batch:
                    emb = sent.get_embedding(self.embeddings["dense"].get_names())
                    if self.similarity_metric == SimilarityMetric.COSINE:
                        emb = emb / torch.linalg.norm(emb)
                    query_embeddings["dense"].append(emb.cpu().numpy())
                    sent.clear_embeddings(self.embeddings["dense"].get_names())

                # Sanity conversion: if flair.device was set as a string, convert to torch.device
                if isinstance(flair.device, str):
                    flair.device = torch.device(flair.device)

                if flair.device.type == "cuda":
                    torch.cuda.empty_cache()

        if self.hybrid_search:
            query_embeddings["sparse"] = []
            self.embeddings["sparse"].embed(inputs)
            for sent in inputs:
                sparse_emb = sent.get_embedding(self.embeddings["sparse"].get_names())
                if self.similarity_metric == SimilarityMetric.COSINE:
                    sparse_emb = sparse_emb / torch.linalg.norm(sparse_emb)
                query_embeddings["sparse"].append(sparse_emb.cpu().numpy())
                sent.clear_embeddings(self.embeddings["sparse"].get_names())

        return {k: np.stack(v, axis=0) for k, v in query_embeddings.items()}
//  [python]
def J_left_SO3(phi):
    phi = phi.view(3, 1)
    phi_norm = torch.norm(phi)
    if torch.abs(phi_norm) > 1e-6:
        a = phi / phi_norm
        J = (torch.sin(phi_norm) / phi_norm) * torch.eye(3, 3, device=phi.device) + \
            (1 - (torch.sin(phi_norm) / phi_norm)) * torch.mm(a, a.transpose(0, 1)) + \
            ((1 - torch.cos(phi_norm)) / phi_norm) * skew3(a)
    else:
        J = torch.eye(3, 3, device=phi.device) + 0.5 * skew3(phi)
    return J
//  [python]
def J_left_SO3(phi):
    phi = phi.view(3, 1)
    phi_norm = torch.linalg.norm(phi)
    if torch.abs(phi_norm) > 1e-6:
        a = phi / phi_norm
        J = (torch.sin(phi_norm) / phi_norm) * torch.eye(3, 3, device=phi.device) + \
            (1 - (torch.sin(phi_norm) / phi_norm)) * torch.mm(a, a.transpose(0, 1)) + \
            ((1 - torch.cos(phi_norm)) / phi_norm) * skew3(a)
    else:
        J = torch.eye(3, 3, device=phi.device) + 0.5 * skew3(phi)
    return J
//  [python]
def J_left_SO3_inv(phi):
    phi = phi.view(3, 1)
    phi_norm = torch.norm(phi)
    if torch.abs(phi_norm) > 1e-6:
        a = phi / phi_norm
        cot_half_phi_norm = 1.0 / torch.tan(phi_norm / 2)
        J_inv = (phi_norm / 2) * cot_half_phi_norm * torch.eye(3, 3, device=phi.device) + \
                (1 - (phi_norm / 2) * cot_half_phi_norm) * \
                torch.mm(a, a.transpose(0, 1)) - (phi_norm / 2) * skew3(a)
    else:
        J_inv = torch.eye(3, 3, device=phi.device) - 0.5 * skew3(phi)
    return J_inv
//  [python]
def J_left_SO3_inv(phi):
    phi = phi.view(3, 1)
    phi_norm = torch.linalg.norm(phi)
    if torch.abs(phi_norm) > 1e-6:
        a = phi / phi_norm
        cot_half_phi_norm = 1.0 / torch.tan(phi_norm / 2)
        J_inv = (phi_norm / 2) * cot_half_phi_norm * torch.eye(3, 3, device=phi.device) + \
                (1 - (phi_norm / 2) * cot_half_phi_norm) * \
                torch.mm(a, a.transpose(0, 1)) - (phi_norm / 2) * skew3(a)
    else:
        J_inv = torch.eye(3, 3, device=phi.device) - 0.5 * skew3(phi)
    return J_inv
//  [python]
def exp_SO3_b(phi):
    """[summary]

    Args:
        phi ([torch.Tensor]): [(B, 3, 1)] 

    Returns:
        [type]: [description]
    """
    eps = 1e-8
    C = torch.zeros(phi.size(0), 3, 3, device=phi.device)

    phi_norm = torch.norm(phi, dim=1, keepdim=True)
    sel = torch.squeeze(phi_norm > eps)

    phi_norm_sel = phi_norm[sel]
    phi_no_sel = phi[~sel]

    if phi_norm_sel.size(0):
        unit_phi_sel = phi[sel] / phi_norm_sel
        unit_phi_skewed_sel = skew3_b(unit_phi_sel)
        C[sel] = torch.eye(3, 3, device=phi.device).repeat([phi_norm_sel.size(0), 1, 1]) + \
                 torch.sin(phi_norm_sel) * unit_phi_skewed_sel + \
                 (1 - torch.cos(phi_norm_sel)) * torch.matmul(unit_phi_skewed_sel, unit_phi_skewed_sel)

    if phi_no_sel.size(0):
        phi_skewed_no_sel = skew3_b(phi_no_sel)
        C[~sel] = torch.eye(3, 3, device=phi.device).repeat([phi_no_sel.size(0), 1, 1]) + phi_skewed_no_sel

    return C
//  [python]
def exp_SO3_b(phi):
    """[summary]

    Args:
        phi ([torch.Tensor]): [(B, 3, 1)] 

    Returns:
        [type]: [description]
    """
    eps = 1e-8
    C = torch.zeros(phi.size(0), 3, 3, device=phi.device)

    phi_norm = torch.linalg.norm(phi, dim=1, keepdim=True)
    sel = torch.squeeze(phi_norm > eps)

    phi_norm_sel = phi_norm[sel]
    phi_no_sel = phi[~sel]

    if phi_norm_sel.size(0):
        unit_phi_sel = phi[sel] / phi_norm_sel
        unit_phi_skewed_sel = skew3_b(unit_phi_sel)
        C[sel] = torch.eye(3, 3, device=phi.device).repeat([phi_norm_sel.size(0), 1, 1]) + \
                 torch.sin(phi_norm_sel) * unit_phi_skewed_sel + \
                 (1 - torch.cos(phi_norm_sel)) * torch.matmul(unit_phi_skewed_sel, unit_phi_skewed_sel)

    if phi_no_sel.size(0):
        phi_skewed_no_sel = skew3_b(phi_no_sel)
        C[~sel] = torch.eye(3, 3, device=phi.device).repeat([phi_no_sel.size(0), 1, 1]) + phi_skewed_no_sel

    return C
//  [python]
def log_fidelity(self, tensors, pd, vd):
        v = tc.einsum('asb,asd->bd',
                      self.tensors[0, :self.vd[0], :self.pd[0], :self.vd[1]],
                      tensors[0, :vd[0], :pd[0], :vd[1]])
        norm = tc.norm(v)
        v = v / norm
        norm = tc.log(norm)
        for n in range(1, self.length):
            if n < self.length - 1:
                v = tc.einsum('ac,asb,csd->bd', v,
                              self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n + 1]],
                              tensors[n, :vd[n], :pd[n], :vd[n + 1]])
            else:
                v = tc.einsum('ac,asb,csd->bd', v,
                              self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n + 1]],
                              tensors[n, :vd[n], :pd[n], :vd[n + 1]])
            _norm = tc.norm(v)
            v = v / _norm
            norm = norm + tc.log(_norm)
        return norm / self.length
//  [python]
def log_fidelity(self, tensors, pd, vd):
        v = tc.einsum('asb,asd->bd',
                      self.tensors[0, :self.vd[0], :self.pd[0], :self.vd[1]],
                      tensors[0, :vd[0], :pd[0], :vd[1]])
        norm = tc.linalg.norm(v)
        v = v / norm
        norm = tc.log(norm)
        for n in range(1, self.length):
            if n < self.length - 1:
                v = tc.einsum('ac,asb,csd->bd', v,
                              self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n + 1]],
                              tensors[n, :vd[n], :pd[n], :vd[n + 1]])
            else:
                v = tc.einsum('ac,asb,csd->bd', v,
                              self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n + 1]],
                              tensors[n, :vd[n], :pd[n], :vd[n + 1]])
            _norm = tc.linalg.norm(v)
            v = v / _norm
            norm = norm + tc.log(_norm)
        return norm / self.length
//  [python]
def J_left_SO3_inv_b(phi):
    eps = 1e-6
    J_inv = torch.zeros(phi.size(0), 3, 3, device=phi.device)
    phi_norm = torch.norm(phi, dim=1, keepdim=True)
    sel = torch.squeeze(phi_norm > eps)

    phi_norm_sel = phi_norm[sel]
    if phi_norm_sel.size(0):
        unit_phi_sel = phi[sel] / phi_norm_sel
        cot_half_phi_norm_sel = 1.0 / torch.tan(phi_norm_sel / 2)
        J_inv[sel] = (phi_norm_sel / 2) * cot_half_phi_norm_sel * \
                     torch.eye(3, 3, device=phi.device).repeat(phi_norm_sel.size(0), 1, 1) + \
                     (1 - (phi_norm_sel / 2) * cot_half_phi_norm_sel) * \
                     torch.matmul(unit_phi_sel, unit_phi_sel.transpose(-2, -1)) - \
                     (phi_norm_sel / 2) * skew3_b(unit_phi_sel)

    phi_no_sel = phi[~sel]
    if phi_no_sel.size(0):
        J_inv[~sel] = torch.eye(3, 3, device=phi.device).repeat(phi_no_sel.size(0), 1, 1) - 0.5 * skew3_b(phi_no_sel)

    return J_inv
//  [python]
def J_left_SO3_inv_b(phi):
    eps = 1e-6
    J_inv = torch.zeros(phi.size(0), 3, 3, device=phi.device)
    phi_norm = torch.linalg.norm(phi, dim=1, keepdim=True)
    sel = torch.squeeze(phi_norm > eps)

    phi_norm_sel = phi_norm[sel]
    if phi_norm_sel.size(0):
        unit_phi_sel = phi[sel] / phi_norm_sel
        cot_half_phi_norm_sel = 1.0 / torch.tan(phi_norm_sel / 2)
        J_inv[sel] = (phi_norm_sel / 2) * cot_half_phi_norm_sel * \
                     torch.eye(3, 3, device=phi.device).repeat(phi_norm_sel.size(0), 1, 1) + \
                     (1 - (phi_norm_sel / 2) * cot_half_phi_norm_sel) * \
                     torch.matmul(unit_phi_sel, unit_phi_sel.transpose(-2, -1)) - \
                     (phi_norm_sel / 2) * skew3_b(unit_phi_sel)

    phi_no_sel = phi[~sel]
    if phi_no_sel.size(0):
        J_inv[~sel] = torch.eye(3, 3, device=phi.device).repeat(phi_no_sel.size(0), 1, 1) - 0.5 * skew3_b(phi_no_sel)

    return J_inv
//  [python]
def test_dist(self, device):
        def run_test(x, y):
            for p in [0, 1, 2, 3, 4, inf, -inf]:
                dist_xy = torch.dist(x, y, p)
                dist_xy_norm = torch.norm(x - y, p)
                self.assertEqual(dist_xy, dist_xy_norm)

        run_test(torch.randn(5, device=device), torch.randn(5, device=device))

        x = torch.zeros(3, device=device)
        y = torch.zeros(3, device=device)
        y[1] = 1.
        run_test(x, y)
//  [python]
def test_dist(self, device):
        def run_test(x, y):
            for p in [0, 1, 2, 3, 4, inf, -inf]:
                dist_xy = torch.dist(x, y, p)
                dist_xy_norm = torch.linalg.norm(x - y, p)
                self.assertEqual(dist_xy, dist_xy_norm)

        run_test(torch.randn(5, device=device), torch.randn(5, device=device))

        x = torch.zeros(3, device=device)
        y = torch.zeros(3, device=device)
        y[1] = 1.
        run_test(x, y)
//  [python]
def _brute_pdist(self, inp, p=2):
        """Computes the same as torch.pdist using primitives"""
        n = inp.shape[-2]
        k = n * (n - 1) // 2
        if k == 0:
            # torch complains about empty indices
            return torch.empty(inp.shape[:-2] + (0,), dtype=inp.dtype, device=inp.device)
        square = torch.norm(inp[..., None, :] - inp[..., None, :, :], p=p, dim=-1)
        unroll = square.view(square.shape[:-2] + (n * n,))
        inds = torch.ones(k, dtype=torch.int)
        inds[torch.arange(n - 1, 1, -1, dtype=torch.int).cumsum(0)] += torch.arange(2, n, dtype=torch.int)
        return unroll[..., inds.cumsum(0)]
//  [python]
def _brute_pdist(self, inp, p=2):
        """Computes the same as torch.pdist using primitives"""
        n = inp.shape[-2]
        k = n * (n - 1) // 2
        if k == 0:
            # torch complains about empty indices
            return torch.empty(inp.shape[:-2] + (0,), dtype=inp.dtype, device=inp.device)
        square = torch.linalg.norm(inp[..., None, :] - inp[..., None, :, :], ord=p, dim=-1)
        unroll = square.view(square.shape[:-2] + (n * n,))
        inds = torch.ones(k, dtype=torch.int)
        inds[torch.arange(n - 1, 1, -1, dtype=torch.int).cumsum(0)] += torch.arange(2, n, dtype=torch.int)
        return unroll[..., inds.cumsum(0)]
//  [python]
def compute_angle(A, B):
    """
     Compute the angle between two tensors of the same size. The tensors will
     be flattened, after which the angle is computed.
    Args:
        A (torch.Tensor): First tensor
        B (torch.Tensor): Second tensor

    Returns: The angle between the two tensors in degrees

    """
    if contains_nan(A):
        print('tensor A contains nans:')
        print(A)
    if contains_nan(B):
        print('tensor B contains nans:')
        print(B)

    inner_product = torch.sum(A*B)  #equal to inner product of flattened tensors
    cosine = inner_product/(torch.norm(A, p='fro')*torch.norm(B, p='fro'))
    if contains_nan(cosine):
        print('cosine contains nans:')
        print('inner product: {}'.format(inner_product))
        print('norm A: {}'.format(torch.norm(A, p='fro')))
        print('norm B: {}'.format(torch.norm(B, p='fro')))

    if cosine > 1 and cosine < 1 + 1e-5:
        cosine = torch.Tensor([1.])
    angle = 180/np.pi*torch.acos(cosine)
    if contains_nan(angle):
        print('angle computation causes NANs. cosines:')
        print(cosine)
    return angle
//  [python]
def compute_angle(A, B):
    """
     Compute the angle between two tensors of the same size. The tensors will
     be flattened, after which the angle is computed.
    Args:
        A (torch.Tensor): First tensor
        B (torch.Tensor): Second tensor

    Returns: The angle between the two tensors in degrees

    """
    if contains_nan(A):
        print('tensor A contains nans:')
        print(A)
    if contains_nan(B):
        print('tensor B contains nans:')
        print(B)

    inner_product = torch.sum(A*B)  #equal to inner product of flattened tensors
    cosine = inner_product/(torch.linalg.norm(A, ord='fro')*torch.linalg.norm(B, ord='fro'))
    if contains_nan(cosine):
        print('cosine contains nans:')
        print('inner product: {}'.format(inner_product))
        print('norm A: {}'.format(torch.linalg.norm(A, ord='fro')))
        print('norm B: {}'.format(torch.linalg.norm(B, ord='fro')))

    if cosine > 1 and cosine < 1 + 1e-5:
        cosine = torch.Tensor([1.])
    angle = 180/np.pi*torch.acos(cosine)
    if contains_nan(angle):
        print('angle computation causes NANs. cosines:')
        print(cosine)
    return angle
//  [python]
def mps_log_norm(self, normalize_mps=False):
        vecs = [None]
        v = tc.einsum('asb,asd->bd',
                      self.tensors[0, :self.vd[0], :self.pd[0], :self.vd[1]],
                      self.tensors[0, :self.vd[0], :self.pd[0], :self.vd[1]])
        norm = tc.norm(v)
        v = v / norm
        vecs.append(v)
        if normalize_mps:
            self.tensors[0, :self.vd[0], :self.pd[0], :self.vd[1]] = \
                self.tensors[0, :self.vd[0], :self.pd[0], :self.vd[1]] / tc.sqrt(norm)
        norm = tc.log(norm)
        for n in range(1, self.length):
            if n < self.length-1:
                v = tc.einsum('ac,asb,csd->bd', v,
                              self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n+1]],
                              self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n+1]])
            else:
                v = tc.einsum('ac,asb,csd->bd', v,
                              self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n + 1]],
                              self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n + 1]])
            _norm = tc.norm(v)
            v = v / _norm
            vecs.append(v)
            norm = norm + tc.log(_norm)
            if normalize_mps:
                self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n+1]] = \
                    self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n+1]] / tc.sqrt(_norm)
        return norm / self.length, vecs
//  [python]
def mps_log_norm(self, normalize_mps=False):
        vecs = [None]
        v = tc.einsum('asb,asd->bd',
                      self.tensors[0, :self.vd[0], :self.pd[0], :self.vd[1]],
                      self.tensors[0, :self.vd[0], :self.pd[0], :self.vd[1]])
        norm = tc.linalg.norm(v)
        v = v / norm
        vecs.append(v)
        if normalize_mps:
            self.tensors[0, :self.vd[0], :self.pd[0], :self.vd[1]] = \
                self.tensors[0, :self.vd[0], :self.pd[0], :self.vd[1]] / tc.sqrt(norm)
        norm = tc.log(norm)
        for n in range(1, self.length):
            if n < self.length-1:
                v = tc.einsum('ac,asb,csd->bd', v,
                              self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n+1]],
                              self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n+1]])
            else:
                v = tc.einsum('ac,asb,csd->bd', v,
                              self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n + 1]],
                              self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n + 1]])
            _norm = tc.linalg.norm(v)
            v = v / _norm
            vecs.append(v)
            norm = norm + tc.log(_norm)
            if normalize_mps:
                self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n+1]] = \
                    self.tensors[n, :self.vd[n], :self.pd[n], :self.vd[n+1]] / tc.sqrt(_norm)
        return norm / self.length, vecs
//  [python]
def J_normalization(x):
    """Jacobian of the L2 normalization, assuming that we normalize
       along the last dimension.
    """
    x_normed = torch.nn.functional.normalize(x, dim=-1)
    norm = torch.norm(x, p=2, dim=-1, keepdim=True)

    Id = torch.diag_embed(torch.ones_like(x_normed))
    J = (Id - x_normed.unsqueeze(-1) @ x_normed.unsqueeze(-2))
    J = J / norm.unsqueeze(-1)
    return J
//  [python]
def J_normalization(x):
    """Jacobian of the L2 normalization, assuming that we normalize
       along the last dimension.
    """
    x_normed = torch.nn.functional.normalize(x, dim=-1)
    norm = torch.linalg.norm(x, ord=2, dim=-1, keepdim=True)

    Id = torch.diag_embed(torch.ones_like(x_normed))
    J = (Id - x_normed.unsqueeze(-1) @ x_normed.unsqueeze(-2))
    J = J / norm.unsqueeze(-1)
    return J
//  [python]
def compute_average_batch_angle(A, B):
    """
    Compute the average of the angles between the mini-batch samples of A and B.
    If the samples of the mini-batch have more than one dimension (minibatch
    dimension not included), the tensors will first be flattened
    Args:
        A (torch.Tensor):  A tensor with as first dimension the mini-batch
            dimension
        B (torch.Tensor): A tensor of the same shape as A

    Returns: The average angle between the two tensors in degrees.

    """

    A = A.flatten(1, -1)
    B = B.flatten(1, -1)
    if contains_nan(A):
        print('tensor A contains nans in activation angles:')
        print(A)
    if contains_nan(B):
        print('tensor B contains nans in activation angles:')
        print(B)
    inner_products = torch.sum(A*B, dim=1)
    A_norms = torch.norm(A, p=2, dim=1)
    B_norms = torch.norm(B, p=2, dim=1)
    cosines = inner_products/(A_norms*B_norms)
    if contains_nan(cosines):
        print('cosines contains nans in activation angles:')
        print('inner product: {}'.format(inner_products))
        print('norm A: {}'.format(A_norms))
        print('norm B: {}'.format(B_norms))
    if torch.sum(A_norms == 0) > 0:
        print('A_norms contains zeros')
    if torch.sum(B_norms == 0) > 0:
        print('B_norms contains zeros')
    cosines = torch.min(cosines, torch.ones_like(cosines))
    angles = torch.acos(cosines)
    return 180/np.pi*torch.mean(angles)
//  [python]
def compute_average_batch_angle(A, B):
    """
    Compute the average of the angles between the mini-batch samples of A and B.
    If the samples of the mini-batch have more than one dimension (minibatch
    dimension not included), the tensors will first be flattened
    Args:
        A (torch.Tensor):  A tensor with as first dimension the mini-batch
            dimension
        B (torch.Tensor): A tensor of the same shape as A

    Returns: The average angle between the two tensors in degrees.

    """

    A = A.flatten(1, -1)
    B = B.flatten(1, -1)
    if contains_nan(A):
        print('tensor A contains nans in activation angles:')
        print(A)
    if contains_nan(B):
        print('tensor B contains nans in activation angles:')
        print(B)
    inner_products = torch.sum(A*B, dim=1)
    A_norms = torch.linalg.norm(A, p=2, dim=1)
    B_norms = torch.linalg.norm(B, p=2, dim=1)
    cosines = inner_products/(A_norms*B_norms)
    if contains_nan(cosines):
        print('cosines contains nans in activation angles:')
        print('inner product: {}'.format(inner_products))
        print('norm A: {}'.format(A_norms))
        print('norm B: {}'.format(B_norms))
    if torch.sum(A_norms == 0) > 0:
        print('A_norms contains zeros')
    if torch.sum(B_norms == 0) > 0:
        print('B_norms contains zeros')
    cosines = torch.min(cosines, torch.ones_like(cosines))
    angles = torch.acos(cosines)
    return 180/np.pi*torch.mean(angles)
//  [python]
def qexp_t(q):
    """
    Applies exponential map to log quaternion
    :param q: N x 3
    :return: N x 4
    """
    n = torch.norm(q, p=2, dim=1, keepdim=True)
    n = torch.clamp(n, min=1e-8)
    q = q * torch.sin(n)
    q = q / n
    q = torch.cat((torch.cos(n), q), dim=1)
    return q
//  [python]
def qexp_t(q):
    """
    Applies exponential map to log quaternion
    :param q: N x 3
    :return: N x 4
    """
    n = torch.linalg.norm(q, ord=2, dim=1, keepdim=True)
    n = torch.clamp(n, min=1e-8)
    q = q * torch.sin(n)
    q = q / n
    q = torch.cat((torch.cos(n), q), dim=1)
    return q
//  [python]
def qlog_t(q):
    """
    Applies the log map to a quaternion
    :param q: N x 4
    :return: N x 3
    """
    n = torch.norm(q[:, 1:], p=2, dim=1, keepdim=True)
    n = torch.clamp(n, min=1e-8)
    q = q[:, 1:] * torch.acos(torch.clamp(q[:, :1], min=-1.0, max=1.0))
    q = q / n
    return q
//  [python]
def qlog_t(q):
    """
    Applies the log map to a quaternion
    :param q: N x 4
    :return: N x 3
    """
    n = torch.linalg.norm(q[:, 1:], ord=2, dim=1, keepdim=True)
    n = torch.clamp(n, min=1e-8)
    q = q[:, 1:] * torch.acos(torch.clamp(q[:, :1], min=-1.0, max=1.0))
    q = q / n
    return q
//  [python]
def norm(x, highdim=False):
    """
    Computes norm of an array of vectors. Given (shape,d), returns (shape) after norm along last dimension
    """
    return torch.norm(x, dim=len(x.shape) - 1)
//  [python]
def norm(x, highdim=False):
    """
    Computes norm of an array of vectors. Given (shape,d), returns (shape) after norm along last dimension
    """
    return torch.linalg.norm(x, dim=len(x.shape) - 1)
//  [python]
def normalize(tensor: torch.Tensor) -> torch.Tensor:
        if tensor.ndim != 2 or tensor.shape[1] != 2:
            raise ValueError("SO2 data tensors can only be 2D vectors.")

        data_norm = torch.norm(tensor, dim=1, keepdim=True)
        near_zero = data_norm < _THESEUS_GLOBAL_PARAMS.get_eps(
            "so2", "norm", tensor.dtype
        )
        data_norm_nz = torch.where(
            near_zero,
            torch.tensor(1.0, dtype=tensor.dtype, device=tensor.device),
            data_norm,
        )
        default_data = torch.tensor(
            [1, 0], dtype=tensor.dtype, device=tensor.device
        ).expand([tensor.shape[0], 2])
        return torch.where(near_zero, default_data, tensor / data_norm_nz)
//  [python]
def normalize(tensor: torch.Tensor) -> torch.Tensor:
        if tensor.ndim != 2 or tensor.shape[1] != 2:
            raise ValueError("SO2 data tensors can only be 2D vectors.")

        data_norm = torch.linalg.norm(tensor, dim=1, keepdim=True)
        near_zero = data_norm < _THESEUS_GLOBAL_PARAMS.get_eps(
            "so2", "norm", tensor.dtype
        )
        data_norm_nz = torch.where(
            near_zero,
            torch.tensor(1.0, dtype=tensor.dtype, device=tensor.device),
            data_norm,
        )
        default_data = torch.tensor(
            [1, 0], dtype=tensor.dtype, device=tensor.device
        ).expand([tensor.shape[0], 2])
        return torch.where(near_zero, default_data, tensor / data_norm_nz)
//  [python]
def __getitem__(self, index):
        # Separate filename and text
        if self.n_speakers > 1:
            audiopath, *extra, text, speaker = self.audiopaths_and_text[index]
            speaker = int(speaker)
        else:
            audiopath, *extra, text = self.audiopaths_and_text[index]
            speaker = None

        pitch = None

        mel = self.get_mel(audiopath)
        text = self.get_text(text)
        energy = torch.norm(mel.float(), dim=0, p=2)
        attn_prior = self.get_prior(index, mel.shape[1], text.shape[0])

        if self.f0_method is not None:
            pitch = self.get_pitch(index, mel.size(-1))
            assert pitch.size(-1) == mel.size(-1)

            # No higher formants?
            if len(pitch.size()) == 1:
                pitch = pitch[None, :]

        return (text, mel, len(text), pitch, energy, speaker, attn_prior,
                audiopath)
//  [python]
def __getitem__(self, index):
        # Separate filename and text
        if self.n_speakers > 1:
            audiopath, *extra, text, speaker = self.audiopaths_and_text[index]
            speaker = int(speaker)
        else:
            audiopath, *extra, text = self.audiopaths_and_text[index]
            speaker = None

        pitch = None

        mel = self.get_mel(audiopath)
        text = self.get_text(text)
        energy = torch.linalg.norm(mel.float(), dim=0, ord=2)
        attn_prior = self.get_prior(index, mel.shape[1], text.shape[0])

        if self.f0_method is not None:
            pitch = self.get_pitch(index, mel.size(-1))
            assert pitch.size(-1) == mel.size(-1)

            # No higher formants?
            if len(pitch.size()) == 1:
                pitch = pitch[None, :]

        return (text, mel, len(text), pitch, energy, speaker, attn_prior,
                audiopath)
//  [python]
def scaled_laplacian(self, node_embeddings, is_eval=False):
        # Normalized graph Laplacian function.
        # :param W: np.ndarray, [n_route, n_route], weighted adjacency matrix of G.
        # :return: np.matrix, [n_route, n_route].
        # learned graph
        node_num = self.num_node
        learned_graph = torch.mm(node_embeddings, node_embeddings.transpose(0, 1))
        norm = torch.norm(node_embeddings, p=2, dim=1, keepdim=True)
        norm = torch.mm(norm, norm.transpose(0, 1))
        learned_graph = learned_graph / norm
        learned_graph = (learned_graph + 1) / 2.
        # learned_graph = F.sigmoid(learned_graph)
        learned_graph = torch.stack([learned_graph, 1-learned_graph], dim=-1)
       
        # make the adj sparse
        if is_eval:
            adj = gumbel_softmax(learned_graph, tau=1, hard=True)
        else:
            adj = gumbel_softmax(learned_graph, tau=1, hard=True)
        adj = adj[:, :, 0].clone().reshape(node_num, -1)
        # mask = torch.eye(self.num_nodes, self.num_nodes).to(device).byte()
        mask = torch.eye(node_num, node_num).bool().cuda()
        adj.masked_fill_(mask, 0)
       
        # d ->  diagonal degree matrix
        W = adj
        n = W.shape[0]
        d = torch.sum(W, axis=1)
        ## L -> graph Laplacian
        L = -W
        L[range(len(L)), range(len(L))] = d
        #for i in range(n):
        #    for j in range(n):
        #        if (d[i] > 0) and (d[j] > 0):
        #            L[i, j] = L[i, j] / torch.sqrt(d[i] * d[j])
        ## lambda_max \approx 2.0, the largest eigenvalues of L.
        try:
            # e, _ = torch.eig(L)
            # lambda_max = e[:, 0].max().detach()
            # import pdb; pdb.set_trace()
            # e = torch.linalg.eigvalsh(L)
            # lambda_max = e.max()
            lambda_max = (L.max() - L.min())
        except Exception as e:
            print("eig error!!: {}".format(e))
            lambda_max = 1.0
       
        # pesudo laplacian matrix, lambda_max = eigs(L.cpu().detach().numpy(), k=1, which='LR')[0][0].real
        tilde = (2 * L / lambda_max - torch.eye(n).cuda())
        self.adj = adj
        self.tilde = tilde
        return adj, tilde
//  [python]
def scaled_laplacian(self, node_embeddings, is_eval=False):
        # Normalized graph Laplacian function.
        # :param W: np.ndarray, [n_route, n_route], weighted adjacency matrix of G.
        # :return: np.matrix, [n_route, n_route].
        # learned graph
        node_num = self.num_node
        learned_graph = torch.mm(node_embeddings, node_embeddings.transpose(0, 1))
        norm = torch.linalg.norm(node_embeddings, p=2, dim=1, keepdim=True)
        norm = torch.mm(norm, norm.transpose(0, 1))
        learned_graph = learned_graph / norm
        learned_graph = (learned_graph + 1) / 2.
        # learned_graph = F.sigmoid(learned_graph)
        learned_graph = torch.stack([learned_graph, 1-learned_graph], dim=-1)
       
        # make the adj sparse
        if is_eval:
            adj = gumbel_softmax(learned_graph, tau=1, hard=True)
        else:
            adj = gumbel_softmax(learned_graph, tau=1, hard=True)
        adj = adj[:, :, 0].clone().reshape(node_num, -1)
        # mask = torch.eye(self.num_nodes, self.num_nodes).to(device).byte()
        mask = torch.eye(node_num, node_num).bool().cuda()
        adj.masked_fill_(mask, 0)
       
        # d ->  diagonal degree matrix
        W = adj
        n = W.shape[0]
        d = torch.sum(W, axis=1)
        ## L -> graph Laplacian
        L = -W
        L[range(len(L)), range(len(L))] = d
        #for i in range(n):
        #    for j in range(n):
        #        if (d[i] > 0) and (d[j] > 0):
        #            L[i, j] = L[i, j] / torch.sqrt(d[i] * d[j])
        ## lambda_max \approx 2.0, the largest eigenvalues of L.
        try:
            # e, _ = torch.eig(L)
            # lambda_max = e[:, 0].max().detach()
            # import pdb; pdb.set_trace()
            # e = torch.linalg.eigvalsh(L)
            # lambda_max = e.max()
            lambda_max = (L.max() - L.min())
        except Exception as e:
            print("eig error!!: {}".format(e))
            lambda_max = 1.0
       
        # pesudo laplacian matrix, lambda_max = eigs(L.cpu().detach().numpy(), k=1, which='LR')[0][0].real
        tilde = (2 * L / lambda_max - torch.eye(n).cuda())
        self.adj = adj
        self.tilde = tilde
        return adj, tilde
//  [python]
def forward(self,
              mean,
              cov=None,
              *,
              decayscales,
              render_color=True):
    """Run MLP. mean is [D, *batch, 3] and cov is [D, *batch, 3, 3]."""
    # NOTE(jainajay): cov is ignored as there IPE for grid embeddings
    # is not implemented.
    batch_shape = mean.shape[:-1]
    num_points = np.prod(batch_shape)
    assert mean.shape[-1] == 3

    x = mean.reshape(num_points, 3)

    # Density network.
    emb_density = neural_fields.mask_grid_features(
        self.emb_density(x), decayscales["density"],
        self.emb_density_config["n_features_per_level"])
    features = self.layers_density(emb_density)
    d, features = features[Ellipsis, :1], features[Ellipsis, 1:]
    assert d.dtype == torch.float32

    # Parameterize canonical frame VolSDF distance as offset from sphere.
    if self.parameterization == "volsdf" and self.base_radius is not None:
      d = d + torch.norm(x, dim=-1, keepdim=True) - self.base_radius

    # Appearance network.
    if render_color:
      emb_color = neural_fields.mask_grid_features(
          self.emb_color(x), decayscales["color"],
          self.emb_color_config["n_features_per_level"])
      features = torch.cat([emb_color, features], dim=-1)
      raw_rgb = self.layers_color(features)
      raw_output = torch.cat([raw_rgb, d], dim=-1)
      return raw_output.view(*batch_shape, 4)

    return d.view(*batch_shape)