jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def Gradient(self, fix_start, fix_end, end_idx=-1):
        self.psi.requires_grad = True
        self.phi.requires_grad = True
        self.omega.requires_grad = True
        self.bond_angles[("C", "N")].requires_grad = True
        self.bond_angles[("N", "CA")].requires_grad = True
        self.bond_angles[("CA", "C")].requires_grad = True
        pred_coords = self.cartesian_coords()
        batch = pred_coords.shape[0]
        coord_reshape = pred_coords.reshape(batch,-1,3,3)
        pred_start = coord_reshape[:,0]
        idx = torch.arange(end_idx.shape[0], device = end_idx.device)
        pred_end = coord_reshape[idx, end_idx]
        
        ## length loss
        length_pred = torch.norm(pred_start.mean(dim=1) - pred_end.mean(dim=1), dim=-1)
        length_true = torch.norm(fix_start.mean(dim=1) - fix_end.mean(dim=1), dim=-1)
        length_loss = torch.mean((length_pred - length_true)**2)

        ## oreintation loss
        def get_coord_system(atoms):
            N = atoms[:,0]
            CA = atoms[:,1]
            C = atoms[:,2]

            unit_vec = lambda x: x / torch.norm(x, dim=-1, keepdim=True)
            u = CA-N
            v = C-CA
            b = unit_vec(u-v)
            n = unit_vec(torch.cross(u,v))
            O = torch.stack((b, n, torch.cross(b, n)), 2) # [batch, coords, axis]
            return O

        O_start_pred = get_coord_system(pred_start)
        O_end_pred = get_coord_system(pred_end)
        O_start_true = get_coord_system(fix_start)
        O_end_true = get_coord_system(fix_end)

        R_pred = torch.matmul(O_start_pred.permute(0,2,1), O_end_pred)
        R_true = torch.matmul(O_start_true.permute(0,2,1), O_end_true)
        q_pred = self._quaternions(R_pred)
        q_true = self._quaternions(R_true)

        ## rotation loss
        rotation_loss = torch.norm(q_pred-q_true, dim=1).mean()

        loss = length_loss + rotation_loss*1000
        loss.backward()
        print("length loss: {}\trot loss: {}".format(length_loss, rotation_loss))

        grad = {"psi": self.psi.grad.clone(),
                "phi": self.phi.grad.clone(),
                "omega": self.omega.grad.clone(),
                "bond_angle_ca_c": self.bond_angles[("CA", "C")].grad.clone(),
                "bond_angle_c_n": self.bond_angles[("C", "N")].grad.clone(),
                "bond_angle_n_ca": self.bond_angles[("N", "CA")].grad.clone()
        }

        self.psi.grad.data.zero_()
        self.phi.grad.data.zero_()
        self.omega.grad.data.zero_()
        self.bond_angles[("CA", "C")].grad.data.zero_()
        self.bond_angles[("C", "N")].grad.data.zero_()
        self.bond_angles[("N", "CA")].grad.data.zero_()

        self.psi.requires_grad = False
        self.phi.requires_grad = False
        self.omega.requires_grad = False
        self.bond_angles[("C", "N")].requires_grad = False
        self.bond_angles[("N", "CA")].requires_grad = False
        self.bond_angles[("CA", "C")].requires_grad = False

        return grad
//  [python]
def Gradient(self, fix_start, fix_end, end_idx=-1):
        self.psi.requires_grad = True
        self.phi.requires_grad = True
        self.omega.requires_grad = True
        self.bond_angles[("C", "N")].requires_grad = True
        self.bond_angles[("N", "CA")].requires_grad = True
        self.bond_angles[("CA", "C")].requires_grad = True
        pred_coords = self.cartesian_coords()
        batch = pred_coords.shape[0]
        coord_reshape = pred_coords.reshape(batch,-1,3,3)
        pred_start = coord_reshape[:,0]
        idx = torch.arange(end_idx.shape[0], device = end_idx.device)
        pred_end = coord_reshape[idx, end_idx]
        
        ## length loss
        length_pred = torch.linalg.norm(pred_start.mean(dim=1) - pred_end.mean(dim=1), dim=-1)
        length_true = torch.linalg.norm(fix_start.mean(dim=1) - fix_end.mean(dim=1), dim=-1)
        length_loss = torch.mean((length_pred - length_true)**2)

        ## oreintation loss
        def get_coord_system(atoms):
            N = atoms[:,0]
            CA = atoms[:,1]
            C = atoms[:,2]

            unit_vec = lambda x: x / torch.linalg.norm(x, dim=-1, keepdim=True)
            u = CA-N
            v = C-CA
            b = unit_vec(u-v)
            n = unit_vec(torch.cross(u,v))
            O = torch.stack((b, n, torch.cross(b, n)), 2) # [batch, coords, axis]
            return O

        O_start_pred = get_coord_system(pred_start)
        O_end_pred = get_coord_system(pred_end)
        O_start_true = get_coord_system(fix_start)
        O_end_true = get_coord_system(fix_end)

        R_pred = torch.matmul(O_start_pred.permute(0,2,1), O_end_pred)
        R_true = torch.matmul(O_start_true.permute(0,2,1), O_end_true)
        q_pred = self._quaternions(R_pred)
        q_true = self._quaternions(R_true)

        ## rotation loss
        rotation_loss = torch.linalg.norm(q_pred-q_true, dim=1).mean()

        loss = length_loss + rotation_loss*1000
        loss.backward()
        print("length loss: {}\trot loss: {}".format(length_loss, rotation_loss))

        grad = {"psi": self.psi.grad.clone(),
                "phi": self.phi.grad.clone(),
                "omega": self.omega.grad.clone(),
                "bond_angle_ca_c": self.bond_angles[("CA", "C")].grad.clone(),
                "bond_angle_c_n": self.bond_angles[("C", "N")].grad.clone(),
                "bond_angle_n_ca": self.bond_angles[("N", "CA")].grad.clone()
        }

        self.psi.grad.data.zero_()
        self.phi.grad.data.zero_()
        self.omega.grad.data.zero_()
        self.bond_angles[("CA", "C")].grad.data.zero_()
        self.bond_angles[("C", "N")].grad.data.zero_()
        self.bond_angles[("N", "CA")].grad.data.zero_()

        self.psi.requires_grad = False
        self.phi.requires_grad = False
        self.omega.requires_grad = False
        self.bond_angles[("C", "N")].requires_grad = False
        self.bond_angles[("N", "CA")].requires_grad = False
        self.bond_angles[("CA", "C")].requires_grad = False

        return grad
//  [python]
def pairwise_distance(self, P0, P1, n_view):
        tensor_view = torch.arange(n_view).cuda()
        combinations_view = torch.combinations(tensor_view)
        distances = torch.zeros(combinations_view.shape[0]).cuda()
        d = (P1 - P0) / torch.norm(P1 - P0, dim=1).unsqueeze(1)  # normalized direction vectors of the cam group

        for comb in range(combinations_view.shape[0]):
            # generate all line direction vectors

            a1 = P0[combinations_view[comb][0],:] #point vector of the first line
            a2 = P0[combinations_view[comb][1],:] #point vector of the second line

            b1 = d[combinations_view[comb][0], :] #direction vector of the first line
            b2 = d[combinations_view[comb][1], :] #direction vector of the second line
            n = torch.cross(b1, b2)
            n_hat = n / torch.norm(n)
            distances[comb] = torch.abs(torch.dot((a1 - a2), n_hat))
        return torch.mean(distances)
//  [python]
def pairwise_distance(self, P0, P1, n_view):
        tensor_view = torch.arange(n_view).cuda()
        combinations_view = torch.combinations(tensor_view)
        distances = torch.zeros(combinations_view.shape[0]).cuda()
        d = (P1 - P0) / torch.linalg.norm(P1 - P0, dim=1, keepdim=True)  # normalized direction vectors of the cam group

        for comb in range(combinations_view.shape[0]):
            # generate all line direction vectors

            a1 = P0[combinations_view[comb][0],:] #point vector of the first line
            a2 = P0[combinations_view[comb][1],:] #point vector of the second line

            b1 = d[combinations_view[comb][0], :] #direction vector of the first line
            b2 = d[combinations_view[comb][1], :] #direction vector of the second line
            n = torch.cross(b1, b2)
            n_hat = n / torch.linalg.norm(n)
            distances[comb] = torch.abs(torch.dot((a1 - a2), n_hat))
        return torch.mean(distances)
//  [python]
def generate_motion(self, load_saved_model=True, samples_to_generate=10, max_steps=300, randomized=True):

        if load_saved_model:
            self.load_best_model()
        self.model.eval()
        test_loader = self.data_loader['test']

        joint_offsets, pos, affs, quat, quat_valid_idx, \
            text, text_valid_idx, intended_emotion, intended_polarity, \
            acting_task, gender, age, handedness, \
            native_tongue = self.return_batch([samples_to_generate], test_loader, randomized=randomized)
        with torch.no_grad():
            joint_lengths = torch.norm(joint_offsets, dim=-1)
            scales, _ = torch.max(joint_lengths, dim=-1)
            quat_pred = torch.zeros_like(quat)
            quat_pred_pre_norm = torch.zeros_like(quat)
            quat_pred[:, 0] = torch.cat(quat_pred.shape[0] * [self.quats_sos]).view(quat_pred[:, 0].shape)
            text_latent = self.model(text, intended_emotion, intended_polarity,
                                     acting_task, gender, age, handedness, native_tongue, only_encoder=True)
            for t in range(1, self.T):
                quat_pred_curr, quat_pred_pre_norm_curr = \
                    self.model(text_latent, quat=quat_pred[:, 0:t],
                               offset_lengths=joint_lengths / scales[..., None],
                               only_decoder=True)
                quat_pred[:, t:t + 1] = quat_pred_curr[:, -1:].clone()
                quat_pred_pre_norm_curr[:, t:t + 1] = quat_pred_pre_norm_curr[:, -1:].clone()

            for s in range(len(quat_pred)):
                quat_pred[s] = qfix(quat_pred[s].view(quat_pred[s].shape[0],
                                                      self.V, -1)).view(quat_pred[s].shape[0], -1)

            root_pos = torch.zeros(quat_pred.shape[0], quat_pred.shape[1], self.C).cuda()
            pos_pred = MocapDataset.forward_kinematics(quat_pred.contiguous().view(
                quat_pred.shape[0], quat_pred.shape[1], -1, self.D), root_pos, self.joint_parents,
                torch.cat((root_pos[:, 0:1], joint_offsets), dim=1).unsqueeze(1))

        animation_pred = {
            'joint_names': self.joint_names,
            'joint_offsets': joint_offsets,
            'joint_parents': self.joint_parents,
            'positions': pos_pred,
            'rotations': quat_pred
        }
        MocapDataset.save_as_bvh(animation_pred,
                                 dataset_name=self.dataset + '_new',
                                 subset_name='test')
        animation = {
            'joint_names': self.joint_names,
            'joint_offsets': joint_offsets,
            'joint_parents': self.joint_parents,
            'positions': pos,
            'rotations': quat
        }
        MocapDataset.save_as_bvh(animation,
                                 dataset_name=self.dataset + '_new',
                                 subset_name='gt')
        pos_pred_np = pos_pred.contiguous().view(pos_pred.shape[0],
                                                 pos_pred.shape[1], -1).permute(0, 2, 1).\
            detach().cpu().numpy()
//  [python]
def generate_motion(self, load_saved_model=True, samples_to_generate=10, max_steps=300, randomized=True):

        if load_saved_model:
            self.load_best_model()
        self.model.eval()
        test_loader = self.data_loader['test']

        joint_offsets, pos, affs, quat, quat_valid_idx, \
            text, text_valid_idx, intended_emotion, intended_polarity, \
            acting_task, gender, age, handedness, \
            native_tongue = self.return_batch([samples_to_generate], test_loader, randomized=randomized)
        with torch.no_grad():
            joint_lengths = torch.linalg.norm(joint_offsets, dim=-1)
            scales, _ = torch.max(joint_lengths, dim=-1)
            quat_pred = torch.zeros_like(quat)
            quat_pred_pre_norm = torch.zeros_like(quat)
            quat_pred[:, 0] = torch.cat(quat_pred.shape[0] * [self.quats_sos]).view(quat_pred[:, 0].shape)
            text_latent = self.model(text, intended_emotion, intended_polarity,
                                     acting_task, gender, age, handedness, native_tongue, only_encoder=True)
            for t in range(1, self.T):
                quat_pred_curr, quat_pred_pre_norm_curr = \
                    self.model(text_latent, quat=quat_pred[:, 0:t],
                               offset_lengths=joint_lengths / scales[..., None],
                               only_decoder=True)
                quat_pred[:, t:t + 1] = quat_pred_curr[:, -1:].clone()
                quat_pred_pre_norm_curr[:, t:t + 1] = quat_pred_pre_norm_curr[:, -1:].clone()

            for s in range(len(quat_pred)):
                quat_pred[s] = qfix(quat_pred[s].view(quat_pred[s].shape[0],
                                                      self.V, -1)).view(quat_pred[s].shape[0], -1)

            root_pos = torch.zeros(quat_pred.shape[0], quat_pred.shape[1], self.C).cuda()
            pos_pred = MocapDataset.forward_kinematics(quat_pred.contiguous().view(
                quat_pred.shape[0], quat_pred.shape[1], -1, self.D), root_pos, self.joint_parents,
                torch.cat((root_pos[:, 0:1], joint_offsets), dim=1).unsqueeze(1))

        animation_pred = {
            'joint_names': self.joint_names,
            'joint_offsets': joint_offsets,
            'joint_parents': self.joint_parents,
            'positions': pos_pred,
            'rotations': quat_pred
        }
        MocapDataset.save_as_bvh(animation_pred,
                                 dataset_name=self.dataset + '_new',
                                 subset_name='test')
        animation = {
            'joint_names': self.joint_names,
            'joint_offsets': joint_offsets,
            'joint_parents': self.joint_parents,
            'positions': pos,
            'rotations': quat
        }
        MocapDataset.save_as_bvh(animation,
                                 dataset_name=self.dataset + '_new',
                                 subset_name='gt')
        pos_pred_np = pos_pred.contiguous().view(pos_pred.shape[0],
                                                 pos_pred.shape[1], -1).permute(0, 2, 1).\
            detach().cpu().numpy()
//  [python]
def __call__(self, batched_verts, batched_verts_gt, batched_uvds, batched_uvds_gt): 
        
        aligned_verts_loss = torch.norm(
            self.align(batched_verts, batched_verts_gt) - batched_verts_gt, dim=2
        ).sum() / (len(batched_verts) * 10)
        
        uvds_loss = torch.norm(
            batched_uvds[:,:,:2] - batched_uvds_gt[:,:,:2], dim=2
        ).mean() * 3

        return dict(
            aligned_verts_loss=aligned_verts_loss, 
            uvds_loss=uvds_loss
        )
//  [python]
def __call__(self, batched_verts, batched_verts_gt, batched_uvds, batched_uvds_gt): 
        
        aligned_verts_loss = torch.linalg.norm(
            self.align(batched_verts, batched_verts_gt) - batched_verts_gt, dim=2
        ).sum() / (len(batched_verts) * 10)
        
        uvds_loss = torch.linalg.norm(
            batched_uvds[:,:,:2] - batched_uvds_gt[:,:,:2], dim=2
        ).mean() * 3

        return dict(
            aligned_verts_loss=aligned_verts_loss, 
            uvds_loss=uvds_loss
        )
//  [python]
def intersect_custom_lstsq(self, P0, P1, ls):
        """P0 and P1 are NxD arrays defining N lines.
        D is the dimension of the space. This function 
        returns the least squares intersection of the N
        lines from the system given by eq. 13 in 
        http://cal.cs.illinois.edu/~johannes/research/LS_line_intersect.pdf.
        """

        # generate all line direction vectors
        n1 = (P1 - P0) / torch.norm(P1 - P0, dim=1).unsqueeze(1)  # normalized

        # generate the array of all projectors
        projs = torch.eye(n1.shape[1]).float().cuda() - n1.unsqueeze(2) * n1.unsqueeze(1)  # I - n*n.T

        # generate R matrix and q vector
        R1 = projs.sum(0)
        q1 = (projs @ P0.unsqueeze(2)).sum(0)

        # solve the least squares problem for the
        # intersection point p: Rp = q
        p1 = ls.lstq(R1, q1)

        return p1
//  [python]
def intersect_custom_lstsq(self, P0, P1, ls):
        """P0 and P1 are NxD arrays defining N lines.
        D is the dimension of the space. This function 
        returns the least squares intersection of the N
        lines from the system given by eq. 13 in 
        http://cal.cs.illinois.edu/~johannes/research/LS_line_intersect.pdf.
        """

        # generate all line direction vectors
        n1 = (P1 - P0) / torch.linalg.norm(P1 - P0, dim=1, keepdim=True)  # normalized

        # generate the array of all projectors
        projs = torch.eye(n1.shape[1]).float().cuda() - n1.unsqueeze(2) * n1.unsqueeze(1)  # I - n*n.T

        # generate R matrix and q vector
        R1 = projs.sum(0)
        q1 = (projs @ P0.unsqueeze(2)).sum(0)

        # solve the least squares problem for the
        # intersection point p: Rp = q
        p1 = ls.lstq(R1, q1)

        return p1
//  [python]
def test_l1_norm(self):
        class NormModel(torch.nn.Module):
            def forward(self, x):
                return torch.norm(x, p=1, dim=-1, keepdim=False)

        x = torch.randn(4, 2, 3, requires_grad=True)
        self.run_test(NormModel(), x)
//  [python]
def test_l1_norm(self):
        class NormModel(torch.nn.Module):
            def forward(self, x):
                return torch.linalg.norm(x, ord=1, dim=-1, keepdim=False)

        x = torch.randn(4, 2, 3, requires_grad=True)
        self.run_test(NormModel(), x)
//  [python]
def test_l2_norm(self):
        class NormModel(torch.nn.Module):
            def forward(self, x):
                return torch.norm(x, p=2, dim=-2, keepdim=False)

        x = torch.randn(4, 2, 3, requires_grad=True)
        self.run_test(NormModel(), x)
//  [python]
def test_l2_norm(self):
        class NormModel(torch.nn.Module):
            def forward(self, x):
                return torch.linalg.norm(x, ord=2, dim=-2, keepdim=False)

        x = torch.randn(4, 2, 3, requires_grad=True)
        self.run_test(NormModel(), x)
//  [python]
def test_dropout_parity(shape, amp, bias, activation, p):
    """
    Check some basic dropout properties
    """

    torch.random.manual_seed(0)
    x = torch.normal(0, 1, size=shape, device="cuda", requires_grad=True)
    b = (
        torch.ones(size=(shape[-1],), device="cuda", requires_grad=True)
        if bias
        else None
    )

    torch.random.manual_seed(0)
    x_ = torch.normal(0, 1, size=shape, device="cuda", requires_grad=True)
    b_ = (
        torch.ones(size=(shape[-1],), device="cuda", requires_grad=True)
        if bias
        else None
    )

    with autocast(enabled=amp):
        torch_activation = build_activation(activation)
        res_torch = torch.nn.functional.dropout(
            torch_activation(x + b if b is not None else x), p=p
        )
        loss_torch = torch.sum(res_torch)

        res_triton = triton_dropout(x=x_, p=p, bias=b_, activation=activation)
        loss_triton = torch.sum(res_triton)

        if p < 0.01:
            # Check the FW pass
            assert torch.allclose(
                loss_torch, loss_triton, rtol=0.01
            ), f"{loss_torch} - {loss_triton}"

            # Check the gradients
            loss_torch.backward()
            loss_triton.backward()

            # - gradients wrt inputs
            assert torch.allclose(
                torch.norm(x.grad), torch.norm(x_.grad), rtol=0.01
            ), f"{x.grad}\n{x_.grad}"

            # - gradients wrt bias
            if bias:
                assert torch.allclose(
                    torch.norm(b.grad), torch.norm(b_.grad), rtol=0.01
                ), f"{b.grad.norm()} - {b_.grad.norm()}"
//  [python]
def test_dropout_parity(shape, amp, bias, activation, p):
    """
    Check some basic dropout properties
    """

    torch.random.manual_seed(0)
    x = torch.normal(0, 1, size=shape, device="cuda", requires_grad=True)
    b = (
        torch.ones(size=(shape[-1],), device="cuda", requires_grad=True)
        if bias
        else None
    )

    torch.random.manual_seed(0)
    x_ = torch.normal(0, 1, size=shape, device="cuda", requires_grad=True)
    b_ = (
        torch.ones(size=(shape[-1],), device="cuda", requires_grad=True)
        if bias
        else None
    )

    with autocast(enabled=amp):
        torch_activation = build_activation(activation)
        res_torch = torch.nn.functional.dropout(
            torch_activation(x + b if b is not None else x), p=p
        )
        loss_torch = torch.sum(res_torch)

        res_triton = triton_dropout(x=x_, p=p, bias=b_, activation=activation)
        loss_triton = torch.sum(res_triton)

        if p < 0.01:
            # Check the FW pass
            assert torch.allclose(
                loss_torch, loss_triton, rtol=0.01
            ), f"{loss_torch} - {loss_triton}"

            # Check the gradients
            loss_torch.backward()
            loss_triton.backward()

            # - gradients wrt inputs
            assert torch.allclose(
                torch.linalg.norm(x.grad), torch.linalg.norm(x_.grad), rtol=0.01
            ), f"{x.grad}\n{x_.grad}"

            # - gradients wrt bias
            if bias:
                assert torch.allclose(
                    torch.linalg.norm(b.grad), torch.linalg.norm(b_.grad), rtol=0.01
                ), f"{b.grad.norm()} - {b_.grad.norm()}"
//  [python]
def test_frobenius_norm_keepdim(self):
        class NormModel(torch.nn.Module):
            def forward(self, x):
                return torch.norm(x, p="fro", dim=(0, 1), keepdim=True)

        x = torch.randn(4, 2, 3, requires_grad=True)
        self.run_test(NormModel(), x)
//  [python]
def test_frobenius_norm_keepdim(self):
        class NormModel(torch.nn.Module):
            def forward(self, x):
                return torch.linalg.norm(x, ord="fro", dim=(0, 1), keepdim=True)

        x = torch.randn(4, 2, 3, requires_grad=True)
        self.run_test(NormModel(), x)
//  [python]
def test_frobenius_norm(self):
        class NormModel(torch.nn.Module):
            def forward(self, x):
                return torch.norm(x, p="fro", dim=0, keepdim=False)

        x = torch.randn(4, 2, 3, requires_grad=True)
        self.run_test(NormModel(), x)
//  [python]
def test_frobenius_norm(self):
        class NormModel(torch.nn.Module):
            def forward(self, x):
                return torch.linalg.norm(x, ord="fro", dim=0, keepdim=False)

        x = torch.randn(4, 2, 3, requires_grad=True)
        self.run_test(NormModel(), x)
//  [python]
def compute_LRA_one(group_xyz, weighting=False):
    B, S, N, C = group_xyz.shape
    dists = torch.norm(group_xyz, dim=-1, keepdim=True) # nn lengths
    
    if weighting:
        dists_max, _ = dists.max(dim=2, keepdim=True)
        dists = dists_max - dists
        dists_sum = dists.sum(dim=2, keepdim=True)
        weights = dists / dists_sum
        weights[weights != weights] = 1.0
        M = torch.matmul(group_xyz.transpose(3,2), weights*group_xyz)
    else:
        M = torch.matmul(group_xyz.transpose(3,2), group_xyz)
    
    eigen_values, vec = M.symeig(eigenvectors=True)
    
    LRA = vec[:,:,:,0]
    LRA_length = torch.norm(LRA, dim=-1, keepdim=True)
    LRA = LRA / LRA_length
    return LRA
//  [python]
def compute_LRA_one(group_xyz, weighting=False):
    B, S, N, C = group_xyz.shape
    dists = torch.linalg.norm(group_xyz, dim=-1, keepdim=True) # nn lengths
    
    if weighting:
        dists_max, _ = dists.max(dim=2, keepdim=True)
        dists = dists_max - dists
        dists_sum = dists.sum(dim=2, keepdim=True)
        weights = dists / dists_sum
        weights[weights != weights] = 1.0
        M = torch.matmul(group_xyz.transpose(3,2), weights*group_xyz)
    else:
        M = torch.matmul(group_xyz.transpose(3,2), group_xyz)
    
    eigen_values, vec = M.symeig(eigenvectors=True)
    
    LRA = vec[:,:,:,0]
    LRA_length = torch.linalg.norm(LRA, dim=-1, keepdim=True)
    LRA = LRA / LRA_length
    return LRA
//  [python]
def compute_LRA(xyz, weighting=False, nsample = 64):
    dists = torch.cdist(xyz, xyz)

    dists, idx = torch.topk(dists, nsample, dim=-1, largest=False, sorted=False)
    dists = dists.unsqueeze(-1)

    group_xyz = index_points(xyz, idx)
    group_xyz = group_xyz - xyz.unsqueeze(2)

    if weighting:
        dists_max, _ = dists.max(dim=2, keepdim=True)
        dists = dists_max - dists
        dists_sum = dists.sum(dim=2, keepdim=True)
        weights = dists / dists_sum
        weights[weights != weights] = 1.0
        M = torch.matmul(group_xyz.transpose(3,2), weights*group_xyz)
    else:
        M = torch.matmul(group_xyz.transpose(3,2), group_xyz)

    eigen_values, vec = M.symeig(eigenvectors=True)

    LRA = vec[:,:,:,0]
    LRA_length = torch.norm(LRA, dim=-1, keepdim=True)
    LRA = LRA / LRA_length
    return LRA
//  [python]
def compute_LRA(xyz, weighting=False, nsample = 64):
    dists = torch.cdist(xyz, xyz)

    dists, idx = torch.topk(dists, nsample, dim=-1, largest=False, sorted=False)
    dists = dists.unsqueeze(-1)

    group_xyz = index_points(xyz, idx)
    group_xyz = group_xyz - xyz.unsqueeze(2)

    if weighting:
        dists_max, _ = dists.max(dim=2, keepdim=True)
        dists = dists_max - dists
        dists_sum = dists.sum(dim=2, keepdim=True)
        weights = dists / dists_sum
        weights[weights != weights] = 1.0
        M = torch.matmul(group_xyz.transpose(3,2), weights*group_xyz)
    else:
        M = torch.matmul(group_xyz.transpose(3,2), group_xyz)

    eigen_values, vec = M.symeig(eigenvectors=True)

    LRA = vec[:,:,:,0]
    LRA_length = torch.linalg.norm(LRA, dim=-1, keepdim=True)
    LRA = LRA / LRA_length
    return LRA
//  [python]
def lanczos_tridiagonalization(hvp_fun, m, v):
    """
    Args:
        hvp_fun: A broadcastable function that computes Hessian-vector products.
        m: number of Lanczos iterations.
        v: (bsz, d) a starting orthonormal vector.
    Returns:
        A tridiagonal matrix (m, m) resulting from the Lanczos method.
    """
    bsz, d = v.shape

    # Multiple torch.stack; need better implementation.
    vecs = [v]
    Q = torch.stack(vecs, dim=1)

    w = hvp_fun(v)
    alpha = torch.einsum("bi,bi->b", w, v)
    w = gram_schmidt_ortho(Q, w)

    alphas = [alpha]
    betas = []

    for j in range(2, m + 1):
        beta = div = torch.norm(w, dim=-1)

        while (div < EPS).any():
            # TODO: only a couple batches are small; more memory efficient method?
            print(f"rerolling {(div < EPS).sum().item()} vectors")
            idx = (beta < EPS)
            w_new = torch.nn.functional.normalize(torch.randn_like(w), dim=-1)
            w_new = gram_schmidt_ortho(Q, w_new)
            w = w * ~idx.unsqueeze(-1) + w_new * idx.unsqueeze(-1)
            div = torch.norm(w, dim=-1)

        v = F.normalize(w, dim=-1)
        vecs.append(v)

        Q = torch.stack(vecs, dim=1)

        w = hvp_fun(v)
        alpha = torch.einsum("bi,bi->b", w, v)
        w = gram_schmidt_ortho(Q, w)

        alphas.append(alpha)
        betas.append(beta)

    alphas = torch.stack(alphas, dim=-1)
    betas = torch.stack(betas, dim=-1)

    T = torch.diag_embed(betas, offset=-1) + torch.diag_embed(alphas, offset=0) + torch.diag_embed(betas, offset=1)
    return T
//  [python]
def lanczos_tridiagonalization(hvp_fun, m, v):
    """
    Args:
        hvp_fun: A broadcastable function that computes Hessian-vector products.
        m: number of Lanczos iterations.
        v: (bsz, d) a starting orthonormal vector.
    Returns:
        A tridiagonal matrix (m, m) resulting from the Lanczos method.
    """
    bsz, d = v.shape

    # Multiple torch.stack; need better implementation.
    vecs = [v]
    Q = torch.stack(vecs, dim=1)

    w = hvp_fun(v)
    alpha = torch.einsum("bi,bi->b", w, v)
    w = gram_schmidt_ortho(Q, w)

    alphas = [alpha]
    betas = []

    for j in range(2, m + 1):
        beta = div = torch.linalg.norm(w, dim=-1)

        while (div < EPS).any():
            # TODO: only a couple batches are small; more memory efficient method?
            print(f"rerolling {(div < EPS).sum().item()} vectors")
            idx = (beta < EPS)
            w_new = torch.nn.functional.normalize(torch.randn_like(w), dim=-1)
            w_new = gram_schmidt_ortho(Q, w_new)
            w = w * ~idx.unsqueeze(-1) + w_new * idx.unsqueeze(-1)
            div = torch.linalg.norm(w, dim=-1)

        v = F.normalize(w, dim=-1)
        vecs.append(v)

        Q = torch.stack(vecs, dim=1)

        w = hvp_fun(v)
        alpha = torch.einsum("bi,bi->b", w, v)
        w = gram_schmidt_ortho(Q, w)

        alphas.append(alpha)
        betas.append(beta)

    alphas = torch.stack(alphas, dim=-1)
    betas = torch.stack(betas, dim=-1)

    T = torch.diag_embed(betas, offset=-1) + torch.diag_embed(alphas, offset=0) + torch.diag_embed(betas, offset=1)
    return T
//  [python]
def order_index(xyz, new_xyz, new_norm, idx):
    B, S, C = new_xyz.shape
    nsample = idx.shape[2]
    grouped_xyz = index_points(xyz, idx)
    grouped_xyz_local = grouped_xyz - new_xyz.view(B, S, 1, C)  # centered

    # project and order
    dist_plane = torch.matmul(grouped_xyz_local, new_norm)
    proj_xyz = grouped_xyz_local - dist_plane*new_norm.view(B, S, 1, C)
    proj_xyz_length = torch.norm(proj_xyz, dim=-1, keepdim=True)
    projected_xyz_unit = proj_xyz / proj_xyz_length
    projected_xyz_unit[projected_xyz_unit != projected_xyz_unit] = 0  # set nan to zero

    length_max_idx = torch.argmax(proj_xyz_length, dim=2)
    vec_ref = projected_xyz_unit.gather(2, length_max_idx.unsqueeze(-1).repeat(1,1,1,3)) # corresponds to the largest length
    
    dots = torch.matmul(projected_xyz_unit, vec_ref.view(B, S, C, 1))
    sign = torch.cross(projected_xyz_unit, vec_ref.view(B, S, 1, C).repeat(1, 1, nsample, 1))
    sign = torch.matmul(sign, new_norm)
    sign = torch.sign(sign)
    sign[:, :, 0, 0] = 1.  # the first is the center point itself, just set sign as 1 to differ from ref_vec 
    dots = sign*dots - (1-sign)
    dots_sorted, indices = torch.sort(dots, dim=2, descending=True)
    idx_ordered = idx.gather(2, indices.squeeze_(-1))

    return dots_sorted, idx_ordered
//  [python]
def order_index(xyz, new_xyz, new_norm, idx):
    B, S, C = new_xyz.shape
    nsample = idx.shape[2]
    grouped_xyz = index_points(xyz, idx)
    grouped_xyz_local = grouped_xyz - new_xyz.view(B, S, 1, C)  # centered

    # project and order
    dist_plane = torch.matmul(grouped_xyz_local, new_norm)
    proj_xyz = grouped_xyz_local - dist_plane*new_norm.view(B, S, 1, C)
    proj_xyz_length = torch.linalg.norm(proj_xyz, dim=-1, keepdim=True)
    projected_xyz_unit = proj_xyz / proj_xyz_length
    projected_xyz_unit[projected_xyz_unit != projected_xyz_unit] = 0  # set nan to zero

    length_max_idx = torch.argmax(proj_xyz_length, dim=2)
    vec_ref = projected_xyz_unit.gather(2, length_max_idx.unsqueeze(-1).repeat(1,1,1,3)) # corresponds to the largest length
    
    dots = torch.matmul(projected_xyz_unit, vec_ref.view(B, S, C, 1))
    sign = torch.cross(projected_xyz_unit, vec_ref.view(B, S, 1, C).repeat(1, 1, nsample, 1))
    sign = torch.matmul(sign, new_norm)
    sign = torch.sign(sign)
    sign[:, :, 0, 0] = 1.  # the first is the center point itself, just set sign as 1 to differ from ref_vec 
    dots = sign*dots - (1-sign)
    dots_sorted, indices = torch.sort(dots, dim=2, descending=True)
    idx_ordered = idx.gather(2, indices.squeeze_(-1))

    return dots_sorted, idx_ordered
//  [python]
def compute_metric(self, v: torch.Tensor, keepdim=False) ->torch.Tensor:
        """Riemannian distance: :math:`d(Z_1, Z_2) = \\sqrt{\\sum_{i=1}^n v_i^2}`.

        Parameters
        ----------
        v : torch.Tensor
            Vector-valued distance
        keepdim : bool
            keep the last dim?

        Returns
        -------
        torch.Tensor
            Riemannian distance between the points
        """
        res = torch.norm(v, dim=-1, keepdim=keepdim)
        return res
//  [python]
def compute_metric(self, v: torch.Tensor, keepdim=False) ->torch.Tensor:
        """Riemannian distance: :math:`d(Z_1, Z_2) = \\sqrt{\\sum_{i=1}^n v_i^2}`.

        Parameters
        ----------
        v : torch.Tensor
            Vector-valued distance
        keepdim : bool
            keep the last dim?

        Returns
        -------
        torch.Tensor
            Riemannian distance between the points
        """
        res = torch.linalg.norm(v, dim=-1, keepdim=keepdim)
        return res
//  [python]
def _log_eucliden_metric(self, x: torch.Tensor, y: torch.Tensor, keepdim=False) ->torch.Tensor:
        """Log-Eucliden Metric distance.

        {}

        References
        ----------
        Log‐Euclidean metrics for fast and simple calculus on diffusion tensors. 2006.
        """.format(self._dist_doc)
        return torch.norm(linalg.sym_logm(x) - linalg.sym_logm(y), dim=[-1, -2], keepdim=keepdim)
//  [python]
def _log_eucliden_metric(self, x: torch.Tensor, y: torch.Tensor, keepdim=False) ->torch.Tensor:
        """Log-Eucliden Metric distance.

        {}

        References
        ----------
        Log‐Euclidean metrics for fast and simple calculus on diffusion tensors. 2006.
        """.format(self._dist_doc)
        return torch.linalg.norm(linalg.sym_logm(x) - linalg.sym_logm(y), dim=[-1, -2], keepdim=keepdim)
//  [python]
def _affine_invariant_metric(self, x: torch.Tensor, y: torch.Tensor, keepdim=False) ->torch.Tensor:
        """Affine Invariant Metric distance.

        {}

        References
        ----------
        A Riemannian framework for tensor computing. 2006.
        """.format(self._dist_doc)
        inv_sqrt_x = linalg.sym_inv_sqrtm1(x)
        return torch.norm(linalg.sym_logm(inv_sqrt_x @ y @ inv_sqrt_x), dim=[-1, -2], keepdim=keepdim)
//  [python]
def _affine_invariant_metric(self, x: torch.Tensor, y: torch.Tensor, keepdim=False) ->torch.Tensor:
        """Affine Invariant Metric distance.

        {}

        References
        ----------
        A Riemannian framework for tensor computing. 2006.
        """.format(self._dist_doc)
        inv_sqrt_x = linalg.sym_inv_sqrtm1(x)
        return torch.linalg.norm(linalg.sym_logm(inv_sqrt_x @ y @ inv_sqrt_x), dim=[-1, -2], keepdim=keepdim)
//  [python]
def get_scale_from_pose(M):
    """
    M [b,4,4]
    return [3]
    """
    rot = M[:, 0:3, 0:3]
    scale = torch.norm(rot, p=2, dim=1) 
    return scale
//  [python]
def get_scale_from_pose(M):
    """
    M [b,4,4]
    return [3]
    """
    rot = M[:, 0:3, 0:3]
    scale = torch.linalg.norm(rot, p=2, dim=1) 
    return scale
//  [python]
def calc_radius(self, batched_data):
        lig_mask = batched_data["lig_mask"]
        pos = batched_data["pos"].clone()
        pos_masked = pos * lig_mask.unsqueeze(-1)
        dist = torch.norm(pos_masked.unsqueeze(1) - pos_masked.unsqueeze(2), dim=-1, p=2) # B x T x T
        max_dist = torch.amax(dist, dim=[1, 2])
        return max_dist
//  [python]
def calc_radius(self, batched_data):
        lig_mask = batched_data["lig_mask"]
        pos = batched_data["pos"].clone()
        pos_masked = pos * lig_mask.unsqueeze(-1)
        dist = torch.linalg.norm(pos_masked.unsqueeze(1) - pos_masked.unsqueeze(2), dim=-1, ord=2) # B x T x T
        max_dist = torch.amax(dist, dim=[1, 2])
        return max_dist
//  [python]
def get_conditioned_derivative(self, batched_data):
        pos = batched_data["pos"] # B x T x 3
        device = pos.device
        n_graphs, n_tokens = pos.size()[:2]
        lig_mask = batched_data["lig_mask"]
        pos_masked = pos * lig_mask.unsqueeze(-1)
        dist = torch.norm(pos_masked.unsqueeze(1) - pos_masked.unsqueeze(2), dim=-1, p=2) # B x T x T
        max_pairs = torch.argmax(dist.reshape(n_graphs, -1), dim=-1)
        max_pair_i, max_pair_j = max_pairs // n_tokens, max_pairs % n_tokens
        max_dist = dist[torch.arange(n_graphs, device=device), max_pair_i, max_pair_j]
        grad_factor = -torch.sign(max_dist - self.target_radius) / (max_dist + 1e-32)
        pos_i = pos[torch.arange(n_graphs, device=device), max_pair_i, :].unsqueeze(1).repeat([1, n_tokens, 1])
        pos_j = pos[torch.arange(n_graphs, device=device), max_pair_j, :].unsqueeze(1).repeat([1, n_tokens, 1])
        mask_i = torch.arange(n_tokens, device=device).unsqueeze(0).repeat([n_graphs, 1]) == max_pair_i.unsqueeze(-1)
        mask_j = torch.arange(n_tokens, device=device).unsqueeze(0).repeat([n_graphs, 1]) == max_pair_j.unsqueeze(-1)
        grad = grad_factor.unsqueeze(1).unsqueeze(2) * (mask_i.unsqueeze(-1) * (pos_i - pos_j) + mask_j.unsqueeze(-1) * (pos_j - pos_i))
        return grad
//  [python]
def get_conditioned_derivative(self, batched_data):
        pos = batched_data["pos"] # B x T x 3
        device = pos.device
        n_graphs, n_tokens = pos.size()[:2]
        lig_mask = batched_data["lig_mask"]
        pos_masked = pos * lig_mask.unsqueeze(-1)
        dist = torch.linalg.norm(pos_masked.unsqueeze(1) - pos_masked.unsqueeze(2), dim=-1, ord=2) # B x T x T
        max_pairs = torch.argmax(dist.reshape(n_graphs, -1), dim=-1)
        max_pair_i, max_pair_j = max_pairs // n_tokens, max_pairs % n_tokens
        max_dist = dist[torch.arange(n_graphs, device=device), max_pair_i, max_pair_j]
        grad_factor = -torch.sign(max_dist - self.target_radius) / (max_dist + 1e-32)
        pos_i = pos[torch.arange(n_graphs, device=device), max_pair_i, :].unsqueeze(1).repeat([1, n_tokens, 1])
        pos_j = pos[torch.arange(n_graphs, device=device), max_pair_j, :].unsqueeze(1).repeat([1, n_tokens, 1])
        mask_i = torch.arange(n_tokens, device=device).unsqueeze(0).repeat([n_graphs, 1]) == max_pair_i.unsqueeze(-1)
        mask_j = torch.arange(n_tokens, device=device).unsqueeze(0).repeat([n_graphs, 1]) == max_pair_j.unsqueeze(-1)
        grad = grad_factor.unsqueeze(1).unsqueeze(2) * (mask_i.unsqueeze(-1) * (pos_i - pos_j) + mask_j.unsqueeze(-1) * (pos_j - pos_i))
        return grad
//  [python]
def axisangle2quat(vec):
    """
    Converts scaled axis-angle to quat.

    Args:
        vec (torch.Tensor): (ax, ay, az) axis-angle exponential coordinates
    Returns:
        torch.Tensor: (x, y, z, w) vec4 float angles
    """
    # Grab angle
    angle = torch.norm(vec)

    # handle zero-rotation case
    if torch.isclose(angle, torch.tensor([0.0]).to(vec.device)):
        return torch.tensor([0.0, 0.0, 0.0, 1.0]).to(vec.device)

    # make sure that axis is a unit vector
    axis = vec / angle

    q = torch.zeros(4, device=vec.device)
    q[3] = torch.cos(angle / 2.0)
    q[:3] = axis * torch.sin(angle / 2.0)
    return q
//  [python]
def axisangle2quat(vec):
    """
    Converts scaled axis-angle to quat.

    Args:
        vec (torch.Tensor): (ax, ay, az) axis-angle exponential coordinates
    Returns:
        torch.Tensor: (x, y, z, w) vec4 float angles
    """
    # Grab angle
    angle = torch.linalg.norm(vec)

    # handle zero-rotation case
    if torch.isclose(angle, torch.tensor([0.0]).to(vec.device)):
        return torch.tensor([0.0, 0.0, 0.0, 1.0]).to(vec.device)

    # make sure that axis is a unit vector
    axis = vec / angle

    q = torch.zeros(4, device=vec.device)
    q[3] = torch.cos(angle / 2.0)
    q[:3] = axis * torch.sin(angle / 2.0)
    return q
//  [python]
def RI_features(xyz, norm, new_xyz, new_norm, idx, group_all=False):
    B, S, C = new_xyz.shape

    new_norm = new_norm.unsqueeze(-1)
    dots_sorted, idx_ordered = order_index(xyz, new_xyz, new_norm, idx)

    epsilon=1e-7
    grouped_xyz = index_points(xyz, idx_ordered)  # [B, npoint, nsample, C]
    if not group_all:
        grouped_xyz_local = grouped_xyz - new_xyz.view(B, S, 1, C)  # centered
    else:
        grouped_xyz_local = grouped_xyz  # treat orgin as center
    grouped_xyz_length = torch.norm(grouped_xyz_local, dim=-1, keepdim=True) # nn lengths
    grouped_xyz_unit = grouped_xyz_local / grouped_xyz_length
    grouped_xyz_unit[grouped_xyz_unit != grouped_xyz_unit] = 0  # set nan to zero
    grouped_xyz_norm = index_points(norm, idx_ordered) # nn neighbor normal vectors
    
    grouped_xyz_angle_0 = torch.matmul(grouped_xyz_unit, new_norm)
    grouped_xyz_angle_1 =  (grouped_xyz_unit * grouped_xyz_norm).sum(-1, keepdim=True)
    grouped_xyz_angle_norm = torch.matmul(grouped_xyz_norm, new_norm)
    grouped_xyz_angle_norm = torch.acos(torch.clamp(grouped_xyz_angle_norm, -1 + epsilon, 1 - epsilon))  #
    D_0 = (grouped_xyz_angle_0 < grouped_xyz_angle_1)
    D_0[D_0 ==0] = -1
    grouped_xyz_angle_norm = D_0.float() * grouped_xyz_angle_norm

    grouped_xyz_inner_vec = grouped_xyz_local - torch.roll(grouped_xyz_local, 1, 2)
    grouped_xyz_inner_length = torch.norm(grouped_xyz_inner_vec, dim=-1, keepdim=True) # nn lengths
    grouped_xyz_inner_unit = grouped_xyz_inner_vec / grouped_xyz_inner_length
    grouped_xyz_inner_unit[grouped_xyz_inner_unit != grouped_xyz_inner_unit] = 0  # set nan to zero
    grouped_xyz_inner_angle_0 = (grouped_xyz_inner_unit * grouped_xyz_norm).sum(-1, keepdim=True)
    grouped_xyz_inner_angle_1 = (grouped_xyz_inner_unit * torch.roll(grouped_xyz_norm, 1, 2)).sum(-1, keepdim=True)
    grouped_xyz_inner_angle_2 = (grouped_xyz_norm * torch.roll(grouped_xyz_norm, 1, 2)).sum(-1, keepdim=True)
    grouped_xyz_inner_angle_2 = torch.acos(torch.clamp(grouped_xyz_inner_angle_2, -1 + epsilon, 1 - epsilon))
    D_1 = (grouped_xyz_inner_angle_0 < grouped_xyz_inner_angle_1)
    D_1[D_1 ==0] = -1
    grouped_xyz_inner_angle_2 = D_1.float() * grouped_xyz_inner_angle_2

    proj_inner_angle_feat = dots_sorted - torch.roll(dots_sorted, 1, 2)
    proj_inner_angle_feat[:,:,0,0] = (-3) - dots_sorted[:,:,-1,0]

    ri_feat = torch.cat([grouped_xyz_length, 
                            proj_inner_angle_feat,
                            grouped_xyz_angle_0,
                            grouped_xyz_angle_1,
                            grouped_xyz_angle_norm,
                            grouped_xyz_inner_angle_0,
                            grouped_xyz_inner_angle_1,
                            grouped_xyz_inner_angle_2], dim=-1)

    return ri_feat, idx_ordered
//  [python]
def RI_features(xyz, norm, new_xyz, new_norm, idx, group_all=False):
    B, S, C = new_xyz.shape

    new_norm = new_norm.unsqueeze(-1)
    dots_sorted, idx_ordered = order_index(xyz, new_xyz, new_norm, idx)

    epsilon=1e-7
    grouped_xyz = index_points(xyz, idx_ordered)  # [B, npoint, nsample, C]
    if not group_all:
        grouped_xyz_local = grouped_xyz - new_xyz.view(B, S, 1, C)  # centered
    else:
        grouped_xyz_local = grouped_xyz  # treat orgin as center
    grouped_xyz_length = torch.linalg.norm(grouped_xyz_local, dim=-1, keepdim=True) # nn lengths
    grouped_xyz_unit = grouped_xyz_local / grouped_xyz_length
    grouped_xyz_unit[grouped_xyz_unit != grouped_xyz_unit] = 0  # set nan to zero
    grouped_xyz_norm = index_points(norm, idx_ordered) # nn neighbor normal vectors
    
    grouped_xyz_angle_0 = torch.matmul(grouped_xyz_unit, new_norm)
    grouped_xyz_angle_1 =  (grouped_xyz_unit * grouped_xyz_norm).sum(-1, keepdim=True)
    grouped_xyz_angle_norm = torch.matmul(grouped_xyz_norm, new_norm)
    grouped_xyz_angle_norm = torch.acos(torch.clamp(grouped_xyz_angle_norm, -1 + epsilon, 1 - epsilon))  #
    D_0 = (grouped_xyz_angle_0 < grouped_xyz_angle_1)
    D_0[D_0 ==0] = -1
    grouped_xyz_angle_norm = D_0.float() * grouped_xyz_angle_norm

    grouped_xyz_inner_vec = grouped_xyz_local - torch.roll(grouped_xyz_local, 1, 2)
    grouped_xyz_inner_length = torch.linalg.norm(grouped_xyz_inner_vec, dim=-1, keepdim=True) # nn lengths
    grouped_xyz_inner_unit = grouped_xyz_inner_vec / grouped_xyz_inner_length
    grouped_xyz_inner_unit[grouped_xyz_inner_unit != grouped_xyz_inner_unit] = 0  # set nan to zero
    grouped_xyz_inner_angle_0 = (grouped_xyz_inner_unit * grouped_xyz_norm).sum(-1, keepdim=True)
    grouped_xyz_inner_angle_1 = (grouped_xyz_inner_unit * torch.roll(grouped_xyz_norm, 1, 2)).sum(-1, keepdim=True)
    grouped_xyz_inner_angle_2 = (grouped_xyz_norm * torch.roll(grouped_xyz_norm, 1, 2)).sum(-1, keepdim=True)
    grouped_xyz_inner_angle_2 = torch.acos(torch.clamp(grouped_xyz_inner_angle_2, -1 + epsilon, 1 - epsilon))
    D_1 = (grouped_xyz_inner_angle_0 < grouped_xyz_inner_angle_1)
    D_1[D_1 ==0] = -1
    grouped_xyz_inner_angle_2 = D_1.float() * grouped_xyz_inner_angle_2

    proj_inner_angle_feat = dots_sorted - torch.roll(dots_sorted, 1, 2)
    proj_inner_angle_feat[:,:,0,0] = (-3) - dots_sorted[:,:,-1,0]

    ri_feat = torch.cat([grouped_xyz_length, 
                            proj_inner_angle_feat,
                            grouped_xyz_angle_0,
                            grouped_xyz_angle_1,
                            grouped_xyz_angle_norm,
                            grouped_xyz_inner_angle_0,
                            grouped_xyz_inner_angle_1,
                            grouped_xyz_inner_angle_2], dim=-1)

    return ri_feat, idx_ordered
//  [python]
def _get_scene_scale(self, sequence_name):
        n = len(self.rotations[sequence_name])

        R = torch.zeros(n, 3, 3)
        T = torch.zeros(n, 3)

        for i, ann in enumerate(self.rotations[sequence_name]):
            R[i, ...] = torch.tensor(self.rotations[sequence_name][i]["R"])
            T[i, ...] = torch.tensor(self.rotations[sequence_name][i]["T"])

        cameras = PerspectiveCameras(R=R, T=T)
        cc = cameras.get_camera_center()
        centeroid = torch.mean(cc, dim=0)
        diff = cc - centeroid

        norm = torch.norm(diff, dim=1)
        scale = torch.max(norm).item()

        return scale
//  [python]
def _get_scene_scale(self, sequence_name):
        n = len(self.rotations[sequence_name])

        R = torch.zeros(n, 3, 3)
        T = torch.zeros(n, 3)

        for i, ann in enumerate(self.rotations[sequence_name]):
            R[i, ...] = torch.tensor(self.rotations[sequence_name][i]["R"])
            T[i, ...] = torch.tensor(self.rotations[sequence_name][i]["T"])

        cameras = PerspectiveCameras(R=R, T=T)
        cc = cameras.get_camera_center()
        centeroid = torch.mean(cc, dim=0)
        diff = cc - centeroid

        norm = torch.linalg.norm(diff, dim=1)
        scale = torch.max(norm).item()

        return scale
//  [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.norm(
        (R.T @ centered_feats.unsqueeze(dim=-1)).squeeze(dim=-1),
        p=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.norm(
        (R.T @ centered_feats.unsqueeze(dim=-1)).squeeze(dim=-1),
        p=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