jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def get_vim_params(model, labels, logits, features):

    # Get classifier weights from model: Conv for Resnet101 and MLP for VGG13
    W = model.state_dict()["module.classifier.weight"].detach().clone().cpu()[...,0,0]
    b = model.state_dict()["module.classifier.bias"].detach().clone().cpu()

    u = -torch.linalg.pinv(W) @ b

    # size of subspace
    # pretty much just a heuristic
    D = 1000 if features.shape[-1] > 1500 else 512
    # Mark: change from features.shape[-1] < 512
    if features.shape[-1] <= 512:
        D = features.shape[-1]//2
    centered_feats = features - u
    U = torch.linalg.eigh(
        centered_feats.T@centered_feats
    ).eigenvectors.flip(-1)
    R = U[:,D:] # eigenvectors in columns
    assert R.shape[0] == features.shape[-1]
    vlogits = torch.linalg.norm(
        (R.T @ centered_feats.unsqueeze(dim=-1)).squeeze(dim=-1),
        ord=2, dim=-1
    )
    alpha = logits.max(dim=-1).values.mean() / vlogits.mean()
    alpha = alpha.item()

    vim_params_dict = {
        "alpha": alpha,
        "u":  u,
        "R": R
    }

    centered_feats = F.normalize(centered_feats, dim=-1)
    U = torch.linalg.eigh(
        centered_feats.T@centered_feats
    ).eigenvectors.flip(-1) # rev order to descending eigenvalue
    R = U[:,D:] # eigenvectors in columns
    assert R.shape[0] == features.shape[-1]
    vlogits = torch.linalg.norm(
        (R.T @ centered_feats.unsqueeze(dim=-1)).squeeze(dim=-1),
        ord=2, dim=-1
    )
    alpha = logits.max(dim=-1).values.mean() / vlogits.mean()
    alpha = alpha.item()
    vim_params_dict.update({
        "norm_alpha": alpha,
        "norm_R": R
    })

    return vim_params_dict
//  [python]
def expand(self, batched_data):
        pos = batched_data["pos"] # B x T x 3
        init_pos = batched_data["init_cell_pos"] # B x T x 3
        atoms = batched_data["x"][:, :, 0]
        batch_size, max_num_atoms = pos.size()[:2]
        cell = batched_data["cell_pred"] # B x 3 x 3
        cell_tensor = torch.tensor(self.cells, device=pos.device).to(cell.dtype).unsqueeze(0).expand(batch_size, -1, -1)
        offset = torch.bmm(cell_tensor, cell) # B x 8 x 3
        expand_pos = pos.unsqueeze(1) + offset.unsqueeze(2) # B x 8 x T x 3
        expand_pos = expand_pos.view(batch_size, -1, 3) # B x (8 x T) x 3
        init_expand_pos = init_pos.unsqueeze(1) + offset.unsqueeze(2) # B x 8 x T x 3
        init_expand_pos = init_expand_pos.view(batch_size, -1, 3) # B x (8 x T) x 3
        expand_dist = torch.norm(pos.unsqueeze(2) - expand_pos.unsqueeze(1), p=2, dim=-1) # B x T x (8 x T)
        expand_mask = ((expand_dist < self.cutoff) & (expand_dist > 1e-5)) # B x T x (8 x T)
        expand_mask = torch.masked_fill(expand_mask, atoms.eq(0).unsqueeze(-1), False)
        expand_mask = (torch.sum(expand_mask, dim=1) > 0) & (~(atoms.eq(0).repeat(1, len(self.cells)))) # B x (8 x T)
        expand_len = torch.sum(expand_mask, dim=-1)
        max_expand_len = torch.max(expand_len)
        outcell_index = torch.zeros([batch_size, max_expand_len], dtype=torch.long, device=pos.device)
        expand_pos_compressed = torch.zeros([batch_size, max_expand_len, 3], dtype=pos.dtype, device=pos.device)
        init_expand_pos_compressed = torch.zeros([batch_size, max_expand_len, 3], dtype=pos.dtype, device=pos.device)
        outcell_all_index = torch.arange(max_num_atoms, dtype=torch.long, device=pos.device).repeat(len(self.cells))
        for i in range(batch_size):
           outcell_index[i, :expand_len[i]] = outcell_all_index[expand_mask[i]]
           expand_pos_compressed[i, :expand_len[i], :] = expand_pos[i, expand_mask[i], :]
           init_expand_pos_compressed[i, :expand_len[i], :] = init_expand_pos[i, expand_mask[i], :]
        batched_data["expand_pos"] = expand_pos_compressed
        batched_data["init_expand_pos"] = init_expand_pos_compressed
        batched_data["expand_len"] = expand_len
        batched_data["outcell_index"] = outcell_index
        batched_data["expand_mask"] = mask_after_k_persample(batch_size, max_expand_len, expand_len)

        return batched_data
//  [python]
def expand(self, batched_data):
        pos = batched_data["pos"] # B x T x 3
        init_pos = batched_data["init_cell_pos"] # B x T x 3
        atoms = batched_data["x"][:, :, 0]
        batch_size, max_num_atoms = pos.size()[:2]
        cell = batched_data["cell_pred"] # B x 3 x 3
        cell_tensor = torch.tensor(self.cells, device=pos.device).to(cell.dtype).unsqueeze(0).expand(batch_size, -1, -1)
        offset = torch.bmm(cell_tensor, cell) # B x 8 x 3
        expand_pos = pos.unsqueeze(1) + offset.unsqueeze(2) # B x 8 x T x 3
        expand_pos = expand_pos.view(batch_size, -1, 3) # B x (8 x T) x 3
        init_expand_pos = init_pos.unsqueeze(1) + offset.unsqueeze(2) # B x 8 x T x 3
        init_expand_pos = init_expand_pos.view(batch_size, -1, 3) # B x (8 x T) x 3
        expand_dist = torch.linalg.norm(pos.unsqueeze(2) - expand_pos.unsqueeze(1), ord=2, dim=-1) # B x T x (8 x T)
        expand_mask = ((expand_dist < self.cutoff) & (expand_dist > 1e-5)) # B x T x (8 x T)
        expand_mask = torch.masked_fill(expand_mask, atoms.eq(0).unsqueeze(-1), False)
        expand_mask = (torch.sum(expand_mask, dim=1) > 0) & (~(atoms.eq(0).repeat(1, len(self.cells)))) # B x (8 x T)
        expand_len = torch.sum(expand_mask, dim=-1)
        max_expand_len = torch.max(expand_len)
        outcell_index = torch.zeros([batch_size, max_expand_len], dtype=torch.long, device=pos.device)
        expand_pos_compressed = torch.zeros([batch_size, max_expand_len, 3], dtype=pos.dtype, device=pos.device)
        init_expand_pos_compressed = torch.zeros([batch_size, max_expand_len, 3], dtype=pos.dtype, device=pos.device)
        outcell_all_index = torch.arange(max_num_atoms, dtype=torch.long, device=pos.device).repeat(len(self.cells))
        for i in range(batch_size):
           outcell_index[i, :expand_len[i]] = outcell_all_index[expand_mask[i]]
           expand_pos_compressed[i, :expand_len[i], :] = expand_pos[i, expand_mask[i], :]
           init_expand_pos_compressed[i, :expand_len[i], :] = init_expand_pos[i, expand_mask[i], :]
        batched_data["expand_pos"] = expand_pos_compressed
        batched_data["init_expand_pos"] = init_expand_pos_compressed
        batched_data["expand_len"] = expand_len
        batched_data["outcell_index"] = outcell_index
        batched_data["expand_mask"] = mask_after_k_persample(batch_size, max_expand_len, expand_len)

        return batched_data
//  [python]
def converged(curr_modules, prev_modules, threshold=1e-4):
    """Test whether modules are converged to a specified threshold.

    Tests for the summed norm of the differences between each set of modules
    being less than the given threshold

    Takes two dictionaries mapping names to modules, the set of names for each dictionary
    should be the same, looping over the set of names, for each name take the difference
    between the associated modules in each dictionary

    """
    if curr_modules.keys() != prev_modules.keys():
        raise ValueError("The keys to the given mappings must have the same set of names of modules")

    summed_norms = torch.tensor(0.)
    if None in prev_modules.values():
        return False
    for name in curr_modules.keys():
        curr_weight = get_module_weight(curr_modules[name])
        prev_weight = get_module_weight(prev_modules[name])

        difference = curr_weight.sub(prev_weight)
        summed_norms += torch.norm(difference)
    return bool(summed_norms < threshold)
//  [python]
def converged(curr_modules, prev_modules, threshold=1e-4):
    """Test whether modules are converged to a specified threshold.

    Tests for the summed norm of the differences between each set of modules
    being less than the given threshold

    Takes two dictionaries mapping names to modules, the set of names for each dictionary
    should be the same, looping over the set of names, for each name take the difference
    between the associated modules in each dictionary

    """
    if curr_modules.keys() != prev_modules.keys():
        raise ValueError("The keys to the given mappings must have the same set of names of modules")

    summed_norms = torch.tensor(0.)
    if None in prev_modules.values():
        return False
    for name in curr_modules.keys():
        curr_weight = get_module_weight(curr_modules[name])
        prev_weight = get_module_weight(prev_modules[name])

        difference = curr_weight.sub(prev_weight)
        summed_norms += torch.linalg.norm(difference)
    return bool(summed_norms < threshold)
//  [python]
def cal_rank_save(self, log):

        tol_thresholds = [0, 1e-5, 1e-4, 1e-3, 0.001, 0.01, 0.1, 0.2]

        log.info("q_comb norm is {}".format(torch.norm(self.q_comb)))
        for tol_thre in tol_thresholds:
            log.info("For tol_thre of {}, q_comb rank is {}".format(tol_thre, torch.linalg.matrix_rank(self.q_comb.float(), tol=tol_thre)))

        if torch.distributed.get_rank() == 0:
            torch.save(self.q_comb, os.path.join(log.path, "q_comb.pth"))
//  [python]
def cal_rank_save(self, log):

        tol_thresholds = [0, 1e-5, 1e-4, 1e-3, 0.001, 0.01, 0.1, 0.2]

        log.info("q_comb norm is {}".format(torch.linalg.norm(self.q_comb)))
        for tol_thre in tol_thresholds:
            log.info("For tol_thre of {}, q_comb rank is {}".format(tol_thre, torch.linalg.matrix_rank(self.q_comb.float(), tol=tol_thre)))

        if torch.distributed.get_rank() == 0:
            torch.save(self.q_comb, os.path.join(log.path, "q_comb.pth"))
//  [python]
def forward(self, h_head, h_tail, rels):
        r"""
        Score triples.

        Parameters
        ----------
        h_head : torch.Tensor
            Head entity features. The tensor is of shape :math:`(E, D)`, where
            :math:`E` is the number of triples, and :math:`D` is the feature size.
        h_tail : torch.Tensor
            Tail entity features. The tensor is of shape :math:`(E, D)`, where
            :math:`E` is the number of triples, and :math:`D` is the feature size.
        rels : torch.Tensor
            Relation types. It is a LongTensor of shape :math:`(E)`, where
            :math:`E` is the number of triples.

        Returns
        -------
        torch.Tensor
            The triple scores. The tensor is of shape :math:`(E)`.
        """
        h_rel = self.rel_emb(rels)
        proj_rel = self.rel_project(rels).reshape(-1, self.nfeats, self.rfeats)
        h_head = (h_head.unsqueeze(1) @ proj_rel).squeeze(1)
        h_tail = (h_tail.unsqueeze(1) @ proj_rel).squeeze(1)

        return -torch.norm(h_head + h_rel - h_tail, p=self.p, dim=-1)
//  [python]
def forward(self, h_head, h_tail, rels):
        r"""
        Score triples.

        Parameters
        ----------
        h_head : torch.Tensor
            Head entity features. The tensor is of shape :math:`(E, D)`, where
            :math:`E` is the number of triples, and :math:`D` is the feature size.
        h_tail : torch.Tensor
            Tail entity features. The tensor is of shape :math:`(E, D)`, where
            :math:`E` is the number of triples, and :math:`D` is the feature size.
        rels : torch.Tensor
            Relation types. It is a LongTensor of shape :math:`(E)`, where
            :math:`E` is the number of triples.

        Returns
        -------
        torch.Tensor
            The triple scores. The tensor is of shape :math:`(E)`.
        """
        h_rel = self.rel_emb(rels)
        proj_rel = self.rel_project(rels).reshape(-1, self.nfeats, self.rfeats)
        h_head = (h_head.unsqueeze(1) @ proj_rel).squeeze(1)
        h_tail = (h_tail.unsqueeze(1) @ proj_rel).squeeze(1)

        return -torch.linalg.norm(h_head + h_rel - h_tail, ord=self.p, dim=-1)
//  [python]
def test_constant_fold_reduceL2(self):
        class ReduceModule(torch.nn.Module):
            def forward(self, x):
                a = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
                b = torch.norm(a, p=2, dim=-2, keepdim=False)
                return b + x

        GLOBALS.export_onnx_opset_version = self.opset_version
        GLOBALS.operator_export_type = OperatorExportTypes.ONNX
        x = torch.ones(2, 3)
        graph, _, __ = self._model_to_graph(
            ReduceModule(), (x,), input_names=["x"], dynamic_axes={"x": [0, 1]}
        )

        for node in graph.nodes():
            self.assertNotEqual(node.kind(), "onnx::ReduceL2")
//  [python]
def test_constant_fold_reduceL2(self):
        class ReduceModule(torch.nn.Module):
            def forward(self, x):
                a = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
                b = torch.linalg.norm(a, ord=2, dim=-2, keepdims=False)
                return b + x

        GLOBALS.export_onnx_opset_version = self.opset_version
        GLOBALS.operator_export_type = OperatorExportTypes.ONNX
        x = torch.ones(2, 3)
        graph, _, __ = self._model_to_graph(
            ReduceModule(), (x,), input_names=["x"], dynamic_axes={"x": [0, 1]}
        )

        for node in graph.nodes():
            self.assertNotEqual(node.kind(), "onnx::ReduceL2")
//  [python]
def test_constant_fold_reduceL1(self):
        class NormModule(torch.nn.Module):
            def forward(self, x):
                a = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
                b = torch.norm(a, p=1, dim=-2)
                return b + x

        GLOBALS.export_onnx_opset_version = self.opset_version
        GLOBALS.operator_export_type = OperatorExportTypes.ONNX
        x = torch.ones(2, 3)
        graph, _, __ = self._model_to_graph(
            NormModule(), (x,), input_names=["x"], dynamic_axes={"x": [0, 1]}
        )

        for node in graph.nodes():
            self.assertNotEqual(node.kind(), "onnx::ReduceL1")
//  [python]
def test_constant_fold_reduceL1(self):
        class NormModule(torch.nn.Module):
            def forward(self, x):
                a = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
                b = torch.linalg.norm(a, ord=1, dim=-2)
                return b + x

        GLOBALS.export_onnx_opset_version = self.opset_version
        GLOBALS.operator_export_type = OperatorExportTypes.ONNX
        x = torch.ones(2, 3)
        graph, _, __ = self._model_to_graph(
            NormModule(), (x,), input_names=["x"], dynamic_axes={"x": [0, 1]}
        )

        for node in graph.nodes():
            self.assertNotEqual(node.kind(), "onnx::ReduceL1")
//  [python]
def idealize_backbone(
    x: Union[AtomTensor, BackboneTensor],
    lr: float = 1e-3,
    n_iter: int = 100,
    inplace: bool = False,
) -> BackboneTensor:
    """Idealizes a protein backbone to more closely resemble idealised geometry.

    Adaptation of an implementation by Sergey Ovchinnikov:
    https://github.com/sokrypton/tf_proteins/blob/master/coord_to_dihedrals_tools.ipynb

    :param x: Tensor representing the backbone (can include sidechain atoms but
        these are not used).
    :type x: Union[AtomTensor, BackboneTensor]
    :param lr: Learning rate to use with the optimiser (Adam), defaults to
        ``1e-3``.
    :type lr: float, optional
    :param n_iter: Number of optimisation steps to make, defaults to ``100``.
    :type n_iter: int, optional
    :return: BackboneTensor with idealised geometry.
    :rtype: BackboneTensor
    """
    if not inplace:
        x = copy.deepcopy(x)
    x_constant = x
    x.requires_grad = True

    def mse(x, y):
        return torch.mean((x - y) ** 2)

    def ideal_loss(x, x_constant) -> torch.Tensor:
        bl = get_backbone_bond_lengths(x)
        ba = get_backbone_bond_angles(x)

        loss_ideal = torch.sum(
            torch.norm(
                torch.tensor(IDEAL_BB_BOND_LENGTHS, device=bl.device) - bl,
                dim=-1,
            )
        ) + torch.sum(
            torch.norm(
                torch.tensor(IDEAL_BB_BOND_ANGLES, device=ba.device) - ba
            ),
            dim=-1,
        )
        loss_ms = mse(x[:, :4, :], x_constant[:, :4, :])

        return loss_ideal + loss_ms

    opt = torch.optim.Adam([x], lr=lr)

    for _ in range(n_iter):
        loss = ideal_loss(x, x_constant)
        opt.zero_grad()
        loss.backward()
        opt.step()
    return x
//  [python]
def idealize_backbone(
    x: Union[AtomTensor, BackboneTensor],
    lr: float = 1e-3,
    n_iter: int = 100,
    inplace: bool = False,
) -> BackboneTensor:
    """Idealizes a protein backbone to more closely resemble idealised geometry.

    Adaptation of an implementation by Sergey Ovchinnikov:
    https://github.com/sokrypton/tf_proteins/blob/master/coord_to_dihedrals_tools.ipynb

    :param x: Tensor representing the backbone (can include sidechain atoms but
        these are not used).
    :type x: Union[AtomTensor, BackboneTensor]
    :param lr: Learning rate to use with the optimiser (Adam), defaults to
        ``1e-3``.
    :type lr: float, optional
    :param n_iter: Number of optimisation steps to make, defaults to ``100``.
    :type n_iter: int, optional
    :return: BackboneTensor with idealised geometry.
    :rtype: BackboneTensor
    """
    if not inplace:
        x = copy.deepcopy(x)
    x_constant = x
    x.requires_grad = True

    def mse(x, y):
        return torch.mean((x - y) ** 2)

    def ideal_loss(x, x_constant) -> torch.Tensor:
        bl = get_backbone_bond_lengths(x)
        ba = get_backbone_bond_angles(x)

        loss_ideal = torch.sum(
            torch.linalg.norm(
                torch.tensor(IDEAL_BB_BOND_LENGTHS, device=bl.device) - bl,
                dim=-1,
            )
        ) + torch.sum(
            torch.linalg.norm(
                torch.tensor(IDEAL_BB_BOND_ANGLES, device=ba.device) - ba
            ),
            dim=-1,
        )
        loss_ms = mse(x[:, :4, :], x_constant[:, :4, :])

        return loss_ideal + loss_ms

    opt = torch.optim.Adam([x], lr=lr)

    for _ in range(n_iter):
        loss = ideal_loss(x, x_constant)
        opt.zero_grad()
        loss.backward()
        opt.step()
    return x
//  [python]
def forward(self, utterances, hidden_init=None):
        """
        Computes the embeddings of a batch of utterance spectrograms.

        :param utterances: batch of mel-scale filterbanks of same duration as a tensor of shape
        (batch_size, n_frames, n_channels)
        :param hidden_init: initial hidden state of the LSTM as a tensor of shape (num_layers,
        batch_size, hidden_size). Will default to a tensor of zeros if None.
        :return: the embeddings as a tensor of shape (batch_size, embedding_size)
        """
        out, (hidden, cell) = self.lstm(utterances, hidden_init)
        embeds_raw = self.relu(self.linear(hidden[-1]))
        embeds = embeds_raw / torch.norm(embeds_raw, dim=1, keepdim=True)
        return embeds
//  [python]
def forward(self, utterances, hidden_init=None):
        """
        Computes the embeddings of a batch of utterance spectrograms.

        :param utterances: batch of mel-scale filterbanks of same duration as a tensor of shape
        (batch_size, n_frames, n_channels)
        :param hidden_init: initial hidden state of the LSTM as a tensor of shape (num_layers,
        batch_size, hidden_size). Will default to a tensor of zeros if None.
        :return: the embeddings as a tensor of shape (batch_size, embedding_size)
        """
        out, (hidden, cell) = self.lstm(utterances, hidden_init)
        embeds_raw = self.relu(self.linear(hidden[-1]))
        embeds = embeds_raw / torch.linalg.norm(embeds_raw, dim=1, keepdim=True)
        return embeds
//  [python]
def forward(self, data: Data) -> Data:
        assert data.pos is not None
        assert data.edge_index is not None
        (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr

        dist = torch.norm(pos[col] - pos[row], p=2, dim=-1).view(-1, 1)

        if self.norm and dist.numel() > 0:
            max_val = float(dist.max()) if self.max is None else self.max

            length = self.interval[1] - self.interval[0]
            dist = length * (dist / max_val) + self.interval[0]

        if pseudo is not None and self.cat:
            pseudo = pseudo.view(-1, 1) if pseudo.dim() == 1 else pseudo
            data.edge_attr = torch.cat([pseudo, dist.type_as(pseudo)], dim=-1)
        else:
            data.edge_attr = dist

        return data
//  [python]
def forward(self, data: Data) -> Data:
        assert data.pos is not None
        assert data.edge_index is not None
        (row, col), pos, pseudo = data.edge_index, data.pos, data.edge_attr

        dist = torch.linalg.norm(pos[col] - pos[row], ord=2, dim=-1).view(-1, 1)

        if self.norm and dist.numel() > 0:
            max_val = float(dist.max()) if self.max is None else self.max

            length = self.interval[1] - self.interval[0]
            dist = length * (dist / max_val) + self.interval[0]

        if pseudo is not None and self.cat:
            pseudo = pseudo.view(-1, 1) if pseudo.dim() == 1 else pseudo
            data.edge_attr = torch.cat([pseudo, dist.type_as(pseudo)], dim=-1)
        else:
            data.edge_attr = dist

        return data
//  [python]
def pose_to_mtx(pose):
    b1 = pose[..., 0:3]
    b1 = b1 / torch.norm(b1, dim=-1)[:, None]
    b2 = pose[..., 3:6] - torch.sum(b1 * pose[..., 3:6], dim=-1)[:, None] * b1
    b2 = b2 / torch.norm(b2, dim=-1)[:, None]
    b3 = torch.cross(b1, b2)

    return torch.stack([b1, b2, b3, pose[..., 6:9]], dim=-1)
//  [python]
def pose_to_mtx(pose):
    b1 = pose[..., 0:3]
    b1 = b1 / torch.linalg.norm(b1, dim=-1)[:, None]
    b2 = pose[..., 3:6] - torch.sum(b1 * pose[..., 3:6], dim=-1)[:, None] * b1
    b2 = b2 / torch.linalg.norm(b2, dim=-1)[:, None]
    b3 = torch.cross(b1, b2)

    return torch.stack([b1, b2, b3, pose[..., 6:9]], dim=-1)
//  [python]
def pseudo_counts(self,
                     src_feats,
                     k=10,
                     kernel_cluster_distance=0.008,
                     kernel_epsilon=0.0001,
                     c=0.001,
                     sm=8):
        counts = np.zeros(shape=(src_feats.size()[0], ))
        for step in range(src_feats.size()[0]):
            ob_dist = torch.norm(src_feats[step] - src_feats, p=2, dim=1)
            ob_dist = torch.sort(ob_dist).values
            ob_dist = ob_dist[:k]
            dist = ob_dist.cpu().numpy()
            # moving average
            dist = dist / np.mean(dist + 1e-11)
            dist = np.max(dist - kernel_cluster_distance, 0)
            kernel = kernel_epsilon / (dist + kernel_epsilon)
            s = np.sqrt(np.sum(kernel)) + c

            if np.isnan(s) or s > sm:
                counts[step] = 0.
            else:
                counts[step] = 1 / s
        return counts
//  [python]
def pseudo_counts(self,
                     src_feats,
                     k=10,
                     kernel_cluster_distance=0.008,
                     kernel_epsilon=0.0001,
                     c=0.001,
                     sm=8):
        counts = np.zeros(shape=(src_feats.size()[0], ))
        for step in range(src_feats.size()[0]):
            ob_dist = torch.linalg.norm(src_feats[step] - src_feats, p=2, dim=1)
            ob_dist = torch.sort(ob_dist).values
            ob_dist = ob_dist[:k]
            dist = ob_dist.cpu().numpy()
            # moving average
            dist = dist / np.mean(dist + 1e-11)
            dist = np.max(dist - kernel_cluster_distance, 0)
            kernel = kernel_epsilon / (dist + kernel_epsilon)
            s = np.sqrt(np.sum(kernel)) + c

            if np.isnan(s) or s > sm:
                counts[step] = 0.
            else:
                counts[step] = 1 / s
        return counts
//  [python]
def R_norm_heatmap(R, n, box_radius):
    x = np.linspace(-box_radius, box_radius, n)
    y = np.linspace(-box_radius, box_radius, n)
    X, Y = np.meshgrid(x, y)
    R_field = np.zeros((n, n))  # values of phi
    for i in range(n):
        for j in range(n):
            pt = torch.tensor([[X[i, j], Y[i, j]]], device=device).float()
            R_field[i, j] = torch.norm(R(pt) ** 2).detach().cpu()
    plt.pcolor(X, Y, R_field)
    plt.colorbar()
//  [python]
def R_norm_heatmap(R, n, box_radius):
    x = np.linspace(-box_radius, box_radius, n)
    y = np.linspace(-box_radius, box_radius, n)
    X, Y = np.meshgrid(x, y)
    R_field = np.zeros((n, n))  # values of phi
    for i in range(n):
        for j in range(n):
            pt = torch.tensor([[X[i, j], Y[i, j]]], device=device).float()
            R_field[i, j] = torch.linalg.norm(R(pt) ** 2).detach().cpu()
    plt.pcolor(X, Y, R_field)
    plt.colorbar()
//  [python]
loha_dim = 16
lora_dim = loha_dim * 2  # same parameter count
weight_rank = 256

# make norm=1 low rank weight for test
test_weight = torch.randn(1280, weight_rank) @ torch.randn(weight_rank, 1280)
test_weight = test_weight / torch.norm(test_weight)
//  [python]
loha_dim = 16
lora_dim = loha_dim * 2  # same parameter count
weight_rank = 256

# make norm=1 low rank weight for test
test_weight = torch.randn(1280, weight_rank) @ torch.randn(weight_rank, 1280)
test_weight = test_weight / torch.linalg.norm(test_weight)
//  [python]
def _cbf_heading(self, state):
        good_heading = state[:2] - self.center
        good_heading_norm = good_heading / torch.norm(good_heading)
        return torch.cos(state[2]) * good_heading_norm[0] + torch.sin(state[2]) * good_heading_norm[1]
//  [python]
def _cbf_heading(self, state):
        good_heading = state[:2] - self.center
        good_heading_norm = good_heading / torch.linalg.norm(good_heading)
        return torch.cos(state[2]) * good_heading_norm[0] + torch.sin(state[2]) * good_heading_norm[1]
//  [python]
def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as axis/angle to quaternions.

    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.

    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
    half_angles = angles * 0.5
    eps = 1e-6
    small_angles = angles.abs() < eps
    large = torch.sin(half_angles) / angles
    small = 0.5 - (angles * angles) / 48
    sin_half_angles_over_angles = torch.where(small_angles, small, large)
    quaternions = torch.cat(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
    )
    return quaternions
//  [python]
def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as axis/angle to quaternions.

    Args:
        axis_angle: Rotations given as a vector in axis angle form,
            as a tensor of shape (..., 3), where the magnitude is
            the angle turned anticlockwise in radians around the
            vector's direction.

    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    angles = torch.linalg.norm(axis_angle, ord=2, dim=-1, keepdim=True)
    half_angles = angles * 0.5
    eps = 1e-6
    small_angles = angles.abs() < eps
    large = torch.sin(half_angles) / angles
    small = 0.5 - (angles * angles) / 48
    sin_half_angles_over_angles = torch.where(small_angles, small, large)
    quaternions = torch.cat(
        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1
    )
    return quaternions
//  [python]
def compute_twist_rotation(rotation_matrix, twist_axis):
    '''
    Compute the twist component of given rotation and twist axis
    https://stackoverflow.com/questions/3684269/component-of-a-quaternion-rotation-around-an-axis
    Parameters
    ----------
    rotation_matrix : Tensor (B, 3, 3,)
        The rotation to convert
    twist_axis : Tensor (B, 3,)
        The twist axis
    Returns
    -------
    Tensor (B, 3, 3)
        The twist rotation
    '''
    quaternion = rotation_matrix_to_quaternion(rotation_matrix)

    twist_axis = twist_axis / (torch.norm(twist_axis, dim=1, keepdim=True) + 1e-9)

    projection = torch.einsum('bi,bi->b', twist_axis, quaternion[:, 1:]).unsqueeze(-1) * twist_axis

    twist_quaternion = torch.cat([quaternion[:, 0:1], projection], dim=1)
    twist_quaternion = twist_quaternion / (torch.norm(twist_quaternion, dim=1, keepdim=True) + 1e-9)

    twist_rotation = quaternion_to_rotation_matrix(twist_quaternion)
    twist_aa = quaternion_to_angle_axis(twist_quaternion)

    twist_angle = torch.sum(twist_aa, dim=1,
                            keepdim=True) / torch.sum(twist_axis, dim=1, keepdim=True)

    return twist_rotation, twist_angle
//  [python]
def compute_twist_rotation(rotation_matrix, twist_axis):
    '''
    Compute the twist component of given rotation and twist axis
    https://stackoverflow.com/questions/3684269/component-of-a-quaternion-rotation-around-an-axis
    Parameters
    ----------
    rotation_matrix : Tensor (B, 3, 3,)
        The rotation to convert
    twist_axis : Tensor (B, 3,)
        The twist axis
    Returns
    -------
    Tensor (B, 3, 3)
        The twist rotation
    '''
    quaternion = rotation_matrix_to_quaternion(rotation_matrix)

    twist_axis = twist_axis / (torch.linalg.norm(twist_axis, dim=1, keepdim=True) + 1e-9)

    projection = torch.einsum('bi,bi->b', twist_axis, quaternion[:, 1:]).unsqueeze(-1) * twist_axis

    twist_quaternion = torch.cat([quaternion[:, 0:1], projection], dim=1)
    twist_quaternion = twist_quaternion / (torch.linalg.norm(twist_quaternion, dim=1, keepdim=True) + 1e-9)

    twist_rotation = quaternion_to_rotation_matrix(twist_quaternion)
    twist_aa = quaternion_to_angle_axis(twist_quaternion)

    twist_angle = torch.sum(twist_aa, dim=1,
                            keepdim=True) / torch.sum(twist_axis, dim=1, keepdim=True)

    return twist_rotation, twist_angle
//  [python]
def get_viewdirs(self, normalize=None, to_world=None, flatten=False, reflect=False, grid=None):
        """Get the view directions of the camera"""

        ones = torch.ones((len(self), 1, *self.hw), dtype=self.dtype, device=self.device)
        rays = self.reconstruct_depth_map(ones, to_world=False, grid=grid)

        if reflect:
            rays[:, 1] = - rays[:, 1]
            rays[:, 2] = - rays[:, 2]

        if normalize is True or normalize == 'unit':
            rays = rays / torch.norm(rays, dim=1).unsqueeze(1)
        if normalize == 'plane':
            rays = rays / torch.norm(rays, dim=1).unsqueeze(1)
            rays = rays / rays[:, [2]]

        if to_world:
            # rays = self.to_world(rays).reshape(len(self), 3, *self.hw)
            rays = self.no_translation().to_world(rays).reshape(len(self), 3, *self.hw)

        if flatten:
            rays = rays.reshape(len(self), 3, -1).permute(0, 2, 1)

        return rays
//  [python]
def get_viewdirs(self, normalize=None, to_world=None, flatten=False, reflect=False, grid=None):
        """Get the view directions of the camera"""

        ones = torch.ones((len(self), 1, *self.hw), dtype=self.dtype, device=self.device)
        rays = self.reconstruct_depth_map(ones, to_world=False, grid=grid)

        if reflect:
            rays[:, 1] = - rays[:, 1]
            rays[:, 2] = - rays[:, 2]

        if normalize is True or normalize == 'unit':
            rays = rays / torch.linalg.norm(rays, dim=1, keepdims=True)
        if normalize == 'plane':
            rays = rays / torch.linalg.norm(rays, dim=1, keepdims=True)
            rays = rays / rays[:, [2]]

        if to_world:
            # rays = self.to_world(rays).reshape(len(self), 3, *self.hw)
            rays = self.no_translation().to_world(rays).reshape(len(self), 3, *self.hw)

        if flatten:
            rays = rays.reshape(len(self), 3, -1).permute(0, 2, 1)

        return rays
//  [python]
def _grad_cbf_heading(self, state):
        """
        α = atan2(y, x)

        [∂ cos(α-θ),     ∂ cos(α - θ)  ∂ cos(θ - α) ]
        [------------,   ------------, -----------  ]
        [∂ x             ∂ y           ∂ θ          ]
        >>> self = ObstacleCBF(torch.rand(2), torch.rand(1))
        >>> x0 = torch.rand(3)
        >>> ajac = self._grad_cbf_heading(x0)
        >>> njac = numerical_jac(lambda x: self._cbf_heading(x), x0, 1e-6)[0]
        >>> testing.assert_allclose(njac, ajac, rtol=1e-3, atol=1e-4)
        """
        gcbf = torch.zeros_like(state)
        θ = state[..., 2]
        good_heading = state[..., :2] - self.center
        ρ = torch.norm(good_heading)
        α = torch.atan2(good_heading[..., 1], good_heading[..., 0])
        # ∂ / ∂ x cos(α-θ) = sin(α - θ) y / ρ²
        gcbf[..., 0] = (α - θ).sin() * good_heading[..., 1] / ρ**2
        # ∂ / ∂ y cos(α-θ) = - sin(α - θ) x / ρ²
        gcbf[..., 1] = - (α - θ).sin() * good_heading[..., 0] / ρ**2
        # ∂ / ∂ θ cos(θ-α) = - sin(θ - α)
        gcbf[..., 2] = - (θ-α).sin()
        return gcbf
//  [python]
def _grad_cbf_heading(self, state):
        """
        α = atan2(y, x)

        [∂ cos(α-θ),     ∂ cos(α - θ)  ∂ cos(θ - α) ]
        [------------,   ------------, -----------  ]
        [∂ x             ∂ y           ∂ θ          ]
        >>> self = ObstacleCBF(torch.rand(2), torch.rand(1))
        >>> x0 = torch.rand(3)
        >>> ajac = self._grad_cbf_heading(x0)
        >>> njac = numerical_jac(lambda x: self._cbf_heading(x), x0, 1e-6)[0]
        >>> testing.assert_allclose(njac, ajac, rtol=1e-3, atol=1e-4)
        """
        gcbf = torch.zeros_like(state)
        θ = state[..., 2]
        good_heading = state[..., :2] - self.center
        ρ = torch.linalg.norm(good_heading)
        α = torch.atan2(good_heading[..., 1], good_heading[..., 0])
        # ∂ / ∂ x cos(α-θ) = sin(α - θ) y / ρ²
        gcbf[..., 0] = (α - θ).sin() * good_heading[..., 1] / ρ**2
        # ∂ / ∂ y cos(α-θ) = - sin(α - θ) x / ρ²
        gcbf[..., 1] = - (α - θ).sin() * good_heading[..., 0] / ρ**2
        # ∂ / ∂ θ cos(θ-α) = - sin(θ - α)
        gcbf[..., 2] = - (θ-α).sin()
        return gcbf
//  [python]
def get_plucker(self):
        """Get the plucker coordinates of the camera"""

        b = len(self)
        ones = torch.ones((b, 1, *self.hw), dtype=self.dtype, device=self.device)
        rays = self.reconstruct_depth_map(ones, to_world=False)
        rays = rays / torch.norm(rays, dim=1).unsqueeze(1)
        orig = self.get_center().view(b, 3, 1, 1).repeat(1, 1, *self.hw)

        orig = orig.view(1, 3, -1).permute(0, 2, 1)
        rays = rays.view(1, 3, -1).permute(0, 2, 1)

        cross = torch.cross(orig, rays, dim=-1)
        plucker = torch.cat((rays, cross), dim=-1)

        return plucker
//  [python]
def get_plucker(self):
        """Get the plucker coordinates of the camera"""

        b = len(self)
        ones = torch.ones((b, 1, *self.hw), dtype=self.dtype, device=self.device)
        rays = self.reconstruct_depth_map(ones, to_world=False)
        rays = rays / torch.linalg.norm(rays, dim=1).unsqueeze(1)
        orig = self.get_center().view(b, 3, 1, 1).repeat(1, 1, *self.hw)

        orig = orig.view(1, 3, -1).permute(0, 2, 1)
        rays = rays.view(1, 3, -1).permute(0, 2, 1)

        cross = torch.cross(orig, rays, dim=-1)
        plucker = torch.cat((rays, cross), dim=-1)

        return plucker
//  [python]
def planczos(self, A, order, x):
        q = x / torch.norm(x, dim=2, keepdim=True)
        V = torch.zeros((x.shape[0], x.shape[1], x.shape[2], order), device=self.device)
        V[:, :, :, 0] = q
        q = q.unsqueeze(-1)
        H = torch.zeros((x.shape[0], x.shape[1], order + 1, order), device=self.device)
        r = A @ q
        H[:, :, 0, 0] = torch.sum(q * r, axis=[-2, -1])

        r = r - H[:, :, 0, 0].unsqueeze(-1).unsqueeze(-1) * q
        H[:, :, 1, 0] = torch.norm(r, dim=2).squeeze(-1)

        for k in range(1, order):
            H[:, :, k - 1, k] = H[:, :, k, k - 1]
            v = q.clone()
            q = r / H[:, :, k - 1, k].unsqueeze(-1).unsqueeze(-1)

            V[:, :, :, k] = q.squeeze(-1)

            r = A @ q
            r = r - H[:, :, 0, 0].unsqueeze(-1).unsqueeze(-1) * v

            H[:, :, k, k] = torch.sum(q * r, axis=[-2, -1])

            r = r - H[:, :, k, k].unsqueeze(-1).unsqueeze(-1) * q
            r = r - V @ (V.permute(0, 1, 3, 2) @ r)
            H[:, :, k + 1, k] = torch.norm(r, dim=2).squeeze(-1)

        return V, H[:, :, :order, :order]
//  [python]
def planczos(self, A, order, x):
        q = x / torch.linalg.norm(x, dim=2, keepdim=True)
        V = torch.zeros((x.shape[0], x.shape[1], x.shape[2], order), device=self.device)
        V[:, :, :, 0] = q
        q = q.unsqueeze(-1)
        H = torch.zeros((x.shape[0], x.shape[1], order + 1, order), device=self.device)
        r = A @ q
        H[:, :, 0, 0] = torch.sum(q * r, axis=[-2, -1])

        r = r - H[:, :, 0, 0].unsqueeze(-1).unsqueeze(-1) * q
        H[:, :, 1, 0] = torch.linalg.norm(r, dim=2).squeeze(-1)

        for k in range(1, order):
            H[:, :, k - 1, k] = H[:, :, k, k - 1]
            v = q.clone()
            q = r / H[:, :, k - 1, k].unsqueeze(-1).unsqueeze(-1)

            V[:, :, :, k] = q.squeeze(-1)

            r = A @ q
            r = r - H[:, :, 0, 0].unsqueeze(-1).unsqueeze(-1) * v

            H[:, :, k, k] = torch.sum(q * r, axis=[-2, -1])

            r = r - H[:, :, k, k].unsqueeze(-1).unsqueeze(-1) * q
            r = r - V @ (V.permute(0, 1, 3, 2) @ r)
            H[:, :, k + 1, k] = torch.linalg.norm(r, dim=2).squeeze(-1)

        return V, H[:, :, :order, :order]
//  [python]
def get_render_rays(self, near=None, far=None, n_rays=None, gt=None):
        """Get the rays for rendering"""

        b = len(self)

        ones = torch.ones((b, 1, *self.hw), dtype=self.dtype, device=self.device)

        rays = self.reconstruct_depth_map(ones, to_world=False)
        rays = rays / torch.norm(rays, dim=1).unsqueeze(1)

        rays[:, 1] = - rays[:, 1]
        rays[:, 2] = - rays[:, 2]

        orig = self.pose[:, :3, -1].view(b, 3, 1, 1).repeat(1, 1, *self.hw)
        rays = self.no_translation().inverted_pose().to_world(rays).reshape(b, 3, *self.hw)

        info = [orig, rays]
        if near is not None:
            info = info + [near * ones]
        if far is not None:
            info = info + [far * ones]
        if gt is not None:
            info = info + [gt]

        rays = torch.cat(info, 1)
        rays = rays.permute(0, 2, 3, 1).reshape(b, -1, rays.shape[1])

        if n_rays is not None:
            idx = torch.randint(0, self.n_pixels, (n_rays,))
            rays = rays[:, idx, :]

        return rays
//  [python]
def get_render_rays(self, near=None, far=None, n_rays=None, gt=None):
        """Get the rays for rendering"""

        b = len(self)

        ones = torch.ones((b, 1, *self.hw), dtype=self.dtype, device=self.device)

        rays = self.reconstruct_depth_map(ones, to_world=False)
        rays = rays / torch.linalg.norm(rays, dim=1).unsqueeze(1)

        rays[:, 1] = - rays[:, 1]
        rays[:, 2] = - rays[:, 2]

        orig = self.pose[:, :3, -1].view(b, 3, 1, 1).repeat(1, 1, *self.hw)
        rays = self.no_translation().inverted_pose().to_world(rays).reshape(b, 3, *self.hw)

        info = [orig, rays]
        if near is not None:
            info = info + [near * ones]
        if far is not None:
            info = info + [far * ones]
        if gt is not None:
            info = info + [gt]

        rays = torch.cat(info, 1)
        rays = rays.permute(0, 2, 3, 1).reshape(b, -1, rays.shape[1])

        if n_rays is not None:
            idx = torch.randint(0, self.n_pixels, (n_rays,))
            rays = rays[:, idx, :]

        return rays
//  [python]
def _orthogonalize(matrix, epsilon=1e-8):
    """
    Applies Gram-Schmidt procedure to orthogonalize a given 2D tensor.
    If epsilon is 0, this is equivalent to `torch.qr(matrix, out=(matrix, _))`,
    but `torch.qr` is very slow, probably because it is not optimized for a matrix that has a small number of columns.
    """
    num_cols = matrix.shape[1]
    for i in range(num_cols):
        # Normalize the i'th column.
        col = matrix[:, i : i + 1]
        # If no epsilon is added here, division by zero may be caused by vanishing gradients.
        # This epsilon is not needed if the input matrix covers the gradients of at least one entire layer in the neural network.
        if epsilon == 0:
            # Note that col ** 2 can underflow/overflow if we use FP16.
            # May need to consider multiplying a scaling factor and dividing it later, or using bfloat16 instead.
            col /= torch.norm(col)
        else:
            col /= torch.norm(col) + epsilon
        # Project it on the rest and remove it.
        if i + 1 < num_cols:
            rest = matrix[:, i + 1 :]
            rest -= torch.sum(col * rest, dim=0) * col
//  [python]
def _orthogonalize(matrix, epsilon=1e-8):
    """
    Applies Gram-Schmidt procedure to orthogonalize a given 2D tensor.
    If epsilon is 0, this is equivalent to `torch.qr(matrix, out=(matrix, _))`,
    but `torch.qr` is very slow, probably because it is not optimized for a matrix that has a small number of columns.
    """
    num_cols = matrix.shape[1]
    for i in range(num_cols):
        # Normalize the i'th column.
        col = matrix[:, i : i + 1]
        # If no epsilon is added here, division by zero may be caused by vanishing gradients.
        # This epsilon is not needed if the input matrix covers the gradients of at least one entire layer in the neural network.
        if epsilon == 0:
            # Note that col ** 2 can underflow/overflow if we use FP16.
            # May need to consider multiplying a scaling factor and dividing it later, or using bfloat16 instead.
            col /= torch.linalg.norm(col)
        else:
            col /= torch.linalg.norm(col) + epsilon
        # Project it on the rest and remove it.
        if i + 1 < num_cols:
            rest = matrix[:, i + 1 :]
            rest -= torch.sum(col * rest, dim=0) * col
//  [python]
def step(self, action):
        current_force = torch.clip(action*self.force_scale, -self.max_force, self.max_force) 
        if (self.delay == 0).all():
            self.force = current_force
        else:
            self.force_history.append(current_force)
            self.force_history.pop(0)
            self.force = torch.zeros((self.env_num, self.dim), device=self.device)
            for i in range(self.delay_max):
                env_mask = (self.delay == i)[:,0]
                self.force[env_mask] = self.force_history[-i-1][env_mask]
        
        self.force += self.disturb
        self.decay_force = self.decay * self.v
        self.force -= self.decay_force
        self.res_dyn_force = self.res_dyn_mlp(torch.cat([self.v, action, self.res_dyn_param], dim=-1)) * self.res_dyn_scale
        self.force += self.res_dyn_force
        
        self.x += self.v * self.tau
        self.v += self.force / self.mass * self.tau
        reward = 1.0 - torch.norm(self.x-self.goal,dim=1) - torch.norm(self.v,dim=1)*0.1
        dist2obstacle = torch.norm(self.x-self.obstacle_pos,dim=1) - self.obstacle_radius
        reward[dist2obstacle < 0] += (dist2obstacle[dist2obstacle < 0]-1.0)

        self.step_cnt += 1

        single_done = self.step_cnt >= self.max_steps
        done = torch.ones(self.env_num, device=self.device)*single_done
        # update disturb
        if self.step_cnt % self.disturb_period == 0:
            self._set_disturb()
        if single_done:
            self.reset()
        info = {
            'mass': self.mass, 
            'disturb': self.disturb,
        }
        return self._get_obs(), reward, done, info
//  [python]
def step(self, action):
        current_force = torch.clip(action*self.force_scale, -self.max_force, self.max_force) 
        if (self.delay == 0).all():
            self.force = current_force
        else:
            self.force_history.append(current_force)
            self.force_history.pop(0)
            self.force = torch.zeros((self.env_num, self.dim), device=self.device)
            for i in range(self.delay_max):
                env_mask = (self.delay == i)[:,0]
                self.force[env_mask] = self.force_history[-i-1][env_mask]
        
        self.force += self.disturb
        self.decay_force = self.decay * self.v
        self.force -= self.decay_force
        self.res_dyn_force = self.res_dyn_mlp(torch.cat([self.v, action, self.res_dyn_param], dim=-1)) * self.res_dyn_scale
        self.force += self.res_dyn_force
        
        self.x += self.v * self.tau
        self.v += self.force / self.mass * self.tau
        reward = 1.0 - torch.linalg.norm(self.x-self.goal,dim=1) - torch.linalg.norm(self.v,dim=1)*0.1
        dist2obstacle = torch.linalg.norm(self.x-self.obstacle_pos,dim=1) - self.obstacle_radius
        reward[dist2obstacle < 0] += (dist2obstacle[dist2obstacle < 0]-1.0)

        self.step_cnt += 1

        single_done = self.step_cnt >= self.max_steps
        done = torch.ones(self.env_num, device=self.device)*single_done
        # update disturb
        if self.step_cnt % self.disturb_period == 0:
            self._set_disturb()
        if single_done:
            self.reset()
        info = {
            'mass': self.mass, 
            'disturb': self.disturb,
        }
        return self._get_obs(), reward, done, info
//  [python]
def get_lddt(input, target, cutoff=15.0):
    """
    The approximate lDDT score, based on AlphaFold2's code
    """
    input_dmat = torch.norm(input[..., None, :] - input, dim=-1, p=2)
    target_dmat = torch.norm(target[..., None, :] - target, dim=-1, p=2)
    dists_mask = (target_dmat < cutoff).float() * (1 - torch.eye(target_dmat.shape[-2]))

    dists_l1 = (target_dmat - input_dmat).abs()
    score = 0.25 * (
        (dists_l1 < 0.5).float()
        + (dists_l1 < 1.0).float()
        + (dists_l1 < 2.0).float()
        + (dists_l1 < 4.0).float()
    )
    norm = 1 / (1e-10 + torch.sum(dists_mask, dim=-1))
    score = norm * (1e-10 + torch.sum(dists_mask * score, dim=-1))
    return score
//  [python]
def get_lddt(input, target, cutoff=15.0):
    """
    The approximate lDDT score, based on AlphaFold2's code
    """
    input_dmat = torch.linalg.norm(input[..., None, :] - input, dim=-1, ord=2)
    target_dmat = torch.linalg.norm(target[..., None, :] - target, dim=-1, ord=2)
    dists_mask = (target_dmat < cutoff).float() * (1 - torch.eye(target_dmat.shape[-2]))

    dists_l1 = (target_dmat - input_dmat).abs()
    score = 0.25 * (
        (dists_l1 < 0.5).float()
        + (dists_l1 < 1.0).float()
        + (dists_l1 < 2.0).float()
        + (dists_l1 < 4.0).float()
    )
    norm = 1 / (1e-10 + torch.sum(dists_mask, dim=-1))
    score = norm * (1e-10 + torch.sum(dists_mask * score, dim=-1))
    return score
//  [python]
def assertTorchAllClose(self, x, y, atol=1e-8, rtol=1e-5, msg=None):
        diff = x.float() - y.float()
        diff_norm = torch.norm(diff)
        other_norm = torch.norm(y.float())

        if msg is None:
            msg = "|input - other| > {} + {} * |other|".format(
                atol, rtol
            )

        self.assertLessEqual(
            diff_norm,
            atol + rtol * other_norm,
            msg=msg,
        )
//  [python]
def assertTorchAllClose(self, x, y, atol=1e-8, rtol=1e-5, msg=None):
        diff = x.float() - y.float()
        diff_norm = torch.linalg.norm(diff)
        other_norm = torch.linalg.norm(y.float())

        if msg is None:
            msg = "|input - other| > {} + {} * |other|".format(
                atol, rtol
            )

        self.assertLessEqual(
            diff_norm,
            atol + rtol * other_norm,
            msg=msg,
        )