jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_forward_norm():
    """test_forward_norm"""
    torch.set_grad_enabled(False)
    input_shape = [1, 3, 10, 10]

    class Norm1(Module):
        def forward(self, *args):
            return torch.norm(args[0], p=float("inf"), dim=None, keepdim=False)

    class Norm2(Module):
        def forward(self, *args):
            return torch.norm(args[0], p=float("-inf"), dim=None, keepdim=False)

    class Norm3(Module):
        def forward(self, *args):
            return torch.norm(args[0], p=float("-inf"), dim=None, keepdim=True)

    class Norm4(Module):
        def forward(self, *args):
            return torch.norm(args[0], p=float("inf"), dim=(1, 2), keepdim=False)

    class Norm5(Module):
        def forward(self, *args):
            return torch.norm(args[0], p=float("inf"), dim=(1), keepdim=True)

    class Norm6(Module):
        def forward(self, *args):
            return torch.norm(args[0], p=float(0.5), dim=(1), keepdim=True)

    class Norm7(Module):
        def forward(self, *args):
            return torch.norm(args[0], p=float(1), dim=None, keepdim=False)

    class Norm8(Module):
        def forward(self, *args):
            return torch.norm(args[0], p=float(2.0), dim=(1), keepdim=True)

    class Norm9(Module):
        def forward(self, *args):
            return torch.norm(args[0], p=float(-0.5), dim=(1, 2), keepdim=True)

    class Norm10(Module):
        def forward(self, *args):
            return torch.norm(args[0], p=float(-2), dim=(1), keepdim=False)

    input_data = torch.rand(input_shape).float()
    verify_model(Norm1().float().eval(), input_data=input_data)
    verify_model(Norm2().float().eval(), input_data=input_data)
    verify_model(Norm3().float().eval(), input_data=input_data)
    verify_model(Norm4().float().eval(), input_data=input_data)
    verify_model(Norm5().float().eval(), input_data=input_data)
    verify_model(Norm6().float().eval(), input_data=input_data)
    verify_model(Norm7().float().eval(), input_data=input_data)
    verify_model(Norm8().float().eval(), input_data=input_data)
    verify_model(Norm9().float().eval(), input_data=input_data)
    verify_model(Norm10().float().eval(), input_data=input_data)
//  [python]
def test_forward_norm():
    """test_forward_norm"""
    torch.set_grad_enabled(False)
    input_shape = [1, 3, 10, 10]

    class Norm1(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0], ord=float("inf"), dim=None, keepdim=False)

    class Norm2(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0], ord=float("-inf"), dim=None, keepdim=False)

    class Norm3(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0], ord=float("-inf"), dim=None, keepdim=True)

    class Norm4(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0], ord=float("inf"), dim=(1, 2), keepdim=False)

    class Norm5(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0], ord=float("inf"), dim=(1), keepdim=True)

    class Norm6(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0], ord=float(0.5), dim=(1), keepdim=True)

    class Norm7(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0], ord=float(1), dim=None, keepdim=False)

    class Norm8(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0], ord=float(2.0), dim=(1), keepdim=True)

    class Norm9(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0], ord=float(-0.5), dim=(1, 2), keepdim=True)

    class Norm10(Module):
        def forward(self, *args):
            return torch.linalg.norm(args[0], ord=float(-2), dim=(1), keepdim=False)

    input_data = torch.rand(input_shape).float()
    verify_model(Norm1().float().eval(), input_data=input_data)
    verify_model(Norm2().float().eval(), input_data=input_data)
    verify_model(Norm3().float().eval(), input_data=input_data)
    verify_model(Norm4().float().eval(), input_data=input_data)
    verify_model(Norm5().float().eval(), input_data=input_data)
    verify_model(Norm6().float().eval(), input_data=input_data)
    verify_model(Norm7().float().eval(), input_data=input_data)
    verify_model(Norm8().float().eval(), input_data=input_data)
    verify_model(Norm9().float().eval(), input_data=input_data)
    verify_model(Norm10().float().eval(), input_data=input_data)
//  [python]
def _optimize(self, alpha: float, lr: float, n_iter: int) -> np.ndarray:
        if math.isclose(alpha, 0.): # no L1 regularization
            return self._optimize_lstsq() # very fast

        ics_ret = torch.from_numpy(self.single_ics[:self.size]).to(self.device)
        ics_mut = torch.from_numpy(self.mutual_ics[:self.size, :self.size]).to(self.device)
        weights = torch.from_numpy(self.weights[:self.size]).to(self.device).requires_grad_()
        optim = torch.optim.Adam([weights], lr=lr)

        loss_ic_min = 1e9 + 7  # An arbitrary big value
        best_weights = weights.cpu().detach().numpy()
        iter_cnt = 0
        for it in count():
            ret_ic_sum = (weights * ics_ret).sum()
            mut_ic_sum = (torch.outer(weights, weights) * ics_mut).sum()
            loss_ic = mut_ic_sum - 2 * ret_ic_sum + 1
            loss_ic_curr = loss_ic.item()

            loss_l1 = torch.norm(weights, p=1)  # type: ignore
            loss = loss_ic + alpha * loss_l1

            optim.zero_grad()
            loss.backward()
            optim.step()

            if loss_ic_min - loss_ic_curr > 1e-6:
                iter_cnt = 0
            else:
                iter_cnt += 1

            if loss_ic_curr < loss_ic_min:
                best_weights = weights.cpu().detach().numpy()
                loss_ic_min = loss_ic_curr

            if iter_cnt >= n_iter or it >= 10000:
                break

        return best_weights
//  [python]
def _optimize(self, alpha: float, lr: float, n_iter: int) -> np.ndarray:
        if math.isclose(alpha, 0.): # no L1 regularization
            return self._optimize_lstsq() # very fast

        ics_ret = torch.from_numpy(self.single_ics[:self.size]).to(self.device)
        ics_mut = torch.from_numpy(self.mutual_ics[:self.size, :self.size]).to(self.device)
        weights = torch.from_numpy(self.weights[:self.size]).to(self.device).requires_grad_()
        optim = torch.optim.Adam([weights], lr=lr)

        loss_ic_min = 1e9 + 7  # An arbitrary big value
        best_weights = weights.cpu().detach().numpy()
        iter_cnt = 0
        for it in count():
            ret_ic_sum = (weights * ics_ret).sum()
            mut_ic_sum = (torch.outer(weights, weights) * ics_mut).sum()
            loss_ic = mut_ic_sum - 2 * ret_ic_sum + 1
            loss_ic_curr = loss_ic.item()

            loss_l1 = torch.linalg.norm(weights, ord=1)  # type: ignore
            loss = loss_ic + alpha * loss_l1

            optim.zero_grad()
            loss.backward()
            optim.step()

            if loss_ic_min - loss_ic_curr > 1e-6:
                iter_cnt = 0
            else:
                iter_cnt += 1

            if loss_ic_curr < loss_ic_min:
                best_weights = weights.cpu().detach().numpy()
                loss_ic_min = loss_ic_curr

            if iter_cnt >= n_iter or it >= 10000:
                break

        return best_weights
//  [python]
def spd_mean_kracher_flow(X : Tensor, G0 : Tensor = None, maxiter : int = 50, dim = 0, weights = None, return_dist = False, return_XT = False) -> Tensor:

    if X.shape[dim] == 1:
        if return_dist:
            return X, torch.tensor([0.0], dtype=X.dtype, device=X.device)
        else:
            return X

    if weights is None:
        n = X.shape[dim]
        weights = torch.ones((*X.shape[:-2], 1, 1), dtype=X.dtype, device=X.device)
        weights /= n

    if G0 is None:
        G = (X * weights).sum(dim=dim, keepdim=True)
    else:
        G = G0.clone()

    nu = 1.
    dist = tau = crit = torch.finfo(X.dtype).max
    i = 0

    while (crit > EPS[X.dtype]) and (i < maxiter) and (nu > EPS[X.dtype]):
        i += 1

        Gsq, Ginvsq = sym_invsqrtm2.apply(G)
        XT = sym_logm.apply(Ginvsq @ X @ Ginvsq)
        GT = (XT * weights).sum(dim=dim, keepdim=True)
        G = Gsq @ sym_expm.apply(nu * GT) @ Gsq

        if return_dist:
            dist = torch.norm(XT - GT, p='fro', dim=(-2,-1))
        crit = torch.norm(GT, p='fro', dim=(-2,-1)).max()
        h = nu * crit
        if h < tau:
            nu = 0.95 * nu
            tau = h
        else:
            nu = 0.5 * nu

    if return_dist:
        return G, dist
    if return_XT:
        return G, XT
    return G
//  [python]
def spd_mean_kracher_flow(X : Tensor, G0 : Tensor = None, maxiter : int = 50, dim = 0, weights = None, return_dist = False, return_XT = False) -> Tensor:

    if X.shape[dim] == 1:
        if return_dist:
            return X, torch.tensor([0.0], dtype=X.dtype, device=X.device)
        else:
            return X

    if weights is None:
        n = X.shape[dim]
        weights = torch.ones((*X.shape[:-2], 1, 1), dtype=X.dtype, device=X.device)
        weights /= n

    if G0 is None:
        G = (X * weights).sum(dim=dim, keepdim=True)
    else:
        G = G0.clone()

    nu = 1.
    dist = tau = crit = torch.finfo(X.dtype).max
    i = 0

    while (crit > EPS[X.dtype]) and (i < maxiter) and (nu > EPS[X.dtype]):
        i += 1

        Gsq, Ginvsq = sym_invsqrtm2.apply(G)
        XT = sym_logm.apply(Ginvsq @ X @ Ginvsq)
        GT = (XT * weights).sum(dim=dim, keepdim=True)
        G = Gsq @ sym_expm.apply(nu * GT) @ Gsq

        if return_dist:
            dist = torch.linalg.norm(XT - GT, p='fro', dim=(-2,-1))
        crit = torch.linalg.norm(GT, p='fro', dim=(-2,-1)).max()
        h = nu * crit
        if h < tau:
            nu = 0.95 * nu
            tau = h
        else:
            nu = 0.5 * nu

    if return_dist:
        return G, dist
    if return_XT:
        return G, XT
    return G
//  [python]
def whitenapply_torch_chunked(X, m, P, dimensions=None):
    if not dimensions:
        dimensions = P.size(0)

    # done in chunks to fit in memory for the whole METU
    for i in range(10):
        low = int(i * (X.size(1) / 10))
        up = int((i+1) * (X.size(1) / 10))
        X[:dimensions, low:up] = torch.mm(P[:dimensions, :], X[:, low:up] - m)
        X[:dimensions, low:up] = X[:dimensions, low:up] / (torch.norm(X[:dimensions, low:up], dim=0, keepdim=True) + 1e-6)
    X = X[:dimensions, :]

    return X
//  [python]
def whitenapply_torch_chunked(X, m, P, dimensions=None):
    if not dimensions:
        dimensions = P.size(0)

    # done in chunks to fit in memory for the whole METU
    for i in range(10):
        low = int(i * (X.size(1) / 10))
        up = int((i+1) * (X.size(1) / 10))
        X[:dimensions, low:up] = torch.mm(P[:dimensions, :], X[:, low:up] - m)
        X[:dimensions, low:up] = X[:dimensions, low:up] / (torch.linalg.norm(X[:dimensions, low:up], dim=0, keepdim=True) + 1e-6)
    X = X[:dimensions, :]

    return X
//  [python]
def is_keyframe(
        self,
        cur_frame_idx,
        last_keyframe_idx,
        cur_frame_visibility_filter,
        occ_aware_visibility,
    ):
        kf_translation = self.config["Training"]["kf_translation"]
        kf_min_translation = self.config["Training"]["kf_min_translation"]
        kf_overlap = self.config["Training"]["kf_overlap"]

        curr_frame = self.cameras[cur_frame_idx]
        last_kf = self.cameras[last_keyframe_idx]
        pose_CW = getWorld2View2(curr_frame.R, curr_frame.T)
        last_kf_CW = getWorld2View2(last_kf.R, last_kf.T)
        last_kf_WC = torch.linalg.inv(last_kf_CW)
        dist = torch.norm((pose_CW @ last_kf_WC)[0:3, 3])
        dist_check = dist > kf_translation * self.median_depth
        dist_check2 = dist > kf_min_translation * self.median_depth

        union = torch.logical_or(
            cur_frame_visibility_filter, occ_aware_visibility[last_keyframe_idx]
        ).count_nonzero()
        intersection = torch.logical_and(
            cur_frame_visibility_filter, occ_aware_visibility[last_keyframe_idx]
        ).count_nonzero()
        point_ratio_2 = intersection / union
        return (point_ratio_2 < kf_overlap and dist_check2) or dist_check
//  [python]
def is_keyframe(
        self,
        cur_frame_idx,
        last_keyframe_idx,
        cur_frame_visibility_filter,
        occ_aware_visibility,
    ):
        kf_translation = self.config["Training"]["kf_translation"]
        kf_min_translation = self.config["Training"]["kf_min_translation"]
        kf_overlap = self.config["Training"]["kf_overlap"]

        curr_frame = self.cameras[cur_frame_idx]
        last_kf = self.cameras[last_keyframe_idx]
        pose_CW = getWorld2View2(curr_frame.R, curr_frame.T)
        last_kf_CW = getWorld2View2(last_kf.R, last_kf.T)
        last_kf_WC = torch.linalg.inv(last_kf_CW)
        dist = torch.linalg.norm((pose_CW @ last_kf_WC)[0:3, 3])
        dist_check = dist > kf_translation * self.median_depth
        dist_check2 = dist > kf_min_translation * self.median_depth

        union = torch.logical_or(
            cur_frame_visibility_filter, occ_aware_visibility[last_keyframe_idx]
        ).count_nonzero()
        intersection = torch.logical_and(
            cur_frame_visibility_filter, occ_aware_visibility[last_keyframe_idx]
        ).count_nonzero()
        point_ratio_2 = intersection / union
        return (point_ratio_2 < kf_overlap and dist_check2) or dist_check
//  [python]
def normalize(vectors):
    if isinstance(vectors, torch.Tensor):
        return vectors / (torch.norm(vectors, dim=-1, keepdim=True) + 1e-10)
    # numpy
    return vectors / (np.linalg.norm(vectors, axis=-1, keepdims=True) + 1e-10)
//  [python]
def normalize(vectors):
    if isinstance(vectors, torch.Tensor):
        return vectors / (torch.linalg.norm(vectors, dim=-1, keepdim=True) + 1e-10)
    # numpy
    return vectors / (np.linalg.norm(vectors, axis=-1, keepdims=True) + 1e-10)
//  [python]
def forward(self, x, mask=None):
        """
        Args:
            x: input features with shape of (num_windows*B, Mh*Mw, C)
            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
        """
        # [batch_size*num_windows, Mh*Mw, total_embed_dim]
        B_, N, C = x.shape
        # qkv(): -> [batch_size*num_windows, Mh*Mw, 3 * total_embed_dim]
        # reshape: -> [batch_size*num_windows, Mh*Mw, 3, num_heads, embed_dim_per_head]
        # permute: -> [3, batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]
        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        # [batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]
        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
        attn = torch.einsum("bhqd, bhkd -> bhqk", q, k) \
               / torch.maximum(torch.norm(q, dim=-1, keepdim=True)
                               * torch.norm(k, dim=-1, keepdim=True).transpose(-2, -1),
                               torch.tensor(1e-06, device=q.device, dtype=q.dtype))
        # transpose: -> [batch_size*num_windows, num_heads, embed_dim_per_head, Mh*Mw]
        # @: multiply -> [batch_size*num_windows, num_heads, Mh*Mw, Mh*Mw]
        # q = q * self.scale
        # cosine -->dot?????? Scaled cosine attention：cosine(q,k)/tau 也许理解的不准确： 控制数值范围有利于训练稳定 （残差块的累加 导致深层难以稳定训练）
        attn /= self.tau.clamp(min=0.01)

        attn = attn + self.get_relative_positional_encodings()
        if mask is not None:
            # mask: [nW, Mh*Mw, Mh*Mw]
            nW = mask.shape[0]  # num_windows
            # attn.view: [batch_size, num_windows, num_heads, Mh*Mw, Mh*Mw]
            # mask.unsqueeze: [1, nW, 1, Mh*Mw, Mh*Mw]
            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
            attn = attn.view(-1, self.num_heads, N, N)
            attn = self.softmax(attn)
        else:
            attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        # @: multiply -> [batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]
        # transpose: -> [batch_size*num_windows, Mh*Mw, num_heads, embed_dim_per_head]
        # reshape: -> [batch_size*num_windows, Mh*Mw, total_embed_dim]
        x = torch.einsum("bhal, bhlv -> bhav", attn, v)
        return x
//  [python]
def forward(self, x, mask=None):
        """
        Args:
            x: input features with shape of (num_windows*B, Mh*Mw, C)
            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
        """
        # [batch_size*num_windows, Mh*Mw, total_embed_dim]
        B_, N, C = x.shape
        # qkv(): -> [batch_size*num_windows, Mh*Mw, 3 * total_embed_dim]
        # reshape: -> [batch_size*num_windows, Mh*Mw, 3, num_heads, embed_dim_per_head]
        # permute: -> [3, batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]
        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        # [batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]
        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
        attn = torch.einsum("bhqd, bhkd -> bhqk", q, k) \
               / torch.maximum(torch.linalg.norm(q, dim=-1, keepdim=True)
                               * torch.linalg.norm(k, dim=-1, keepdim=True).transpose(-2, -1),
                               torch.tensor(1e-06, device=q.device, dtype=q.dtype))
        # transpose: -> [batch_size*num_windows, num_heads, embed_dim_per_head, Mh*Mw]
        # @: multiply -> [batch_size*num_windows, num_heads, Mh*Mw, Mh*Mw]
        # q = q * self.scale
        # cosine -->dot?????? Scaled cosine attention：cosine(q,k)/tau 也许理解的不准确： 控制数值范围有利于训练稳定 （残差块的累加 导致深层难以稳定训练）
        attn /= self.tau.clamp(min=0.01)

        attn = attn + self.get_relative_positional_encodings()
        if mask is not None:
            # mask: [nW, Mh*Mw, Mh*Mw]
            nW = mask.shape[0]  # num_windows
            # attn.view: [batch_size, num_windows, num_heads, Mh*Mw, Mh*Mw]
            # mask.unsqueeze: [1, nW, 1, Mh*Mw, Mh*Mw]
            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
            attn = attn.view(-1, self.num_heads, N, N)
            attn = self.softmax(attn)
        else:
            attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        # @: multiply -> [batch_size*num_windows, num_heads, Mh*Mw, embed_dim_per_head]
        # transpose: -> [batch_size*num_windows, Mh*Mw, num_heads, embed_dim_per_head]
        # reshape: -> [batch_size*num_windows, Mh*Mw, total_embed_dim]
        x = torch.einsum("bhal, bhlv -> bhav", attn, v)
        return x
//  [python]
def compute_world2aligned_joints_mat(joints):
    '''
    Compute world to canonical frame (rotation around up axis)
    from the given batch of joints (B x J x 3)
    '''
    left_idx = SMPL_JOINTS['leftUpLeg']
    right_idx = SMPL_JOINTS['rightUpLeg']

    body_right = joints[:, right_idx] - joints[:, left_idx]
    body_right = body_right / torch.norm(body_right, dim=1, keepdim=True)

    world2aligned_mat, world2aligned_aa = compute_aligned_from_right(body_right)

    return world2aligned_mat
//  [python]
def compute_world2aligned_joints_mat(joints):
    '''
    Compute world to canonical frame (rotation around up axis)
    from the given batch of joints (B x J x 3)
    '''
    left_idx = SMPL_JOINTS['leftUpLeg']
    right_idx = SMPL_JOINTS['rightUpLeg']

    body_right = joints[:, right_idx] - joints[:, left_idx]
    body_right = body_right / torch.linalg.norm(body_right, dim=1, keepdim=True)

    world2aligned_mat, world2aligned_aa = compute_aligned_from_right(body_right)

    return world2aligned_mat
//  [python]
def newton_schultz(A, sol, maxiter=100, tol=1e-4, eps=0.1):
    """Matrix square root and inverse square root for SPD matrices"""
    d = len(A)
    normA = torch.norm(A)
    Y = A / ((1 + eps) * normA)
    Id = torch.eye(d, dtype=A.dtype, device=A.device)
    normId = torch.norm(Id)
    ratio = (normA / normId) ** 0.5

    Z = Id / ((1 + eps) * normId)
    for ii in range(maxiter):
        T = 0.5 * (3 * Id - Z @ Y)
        Y = Y @ T
        Z = T @ Z
        if abs(sol - Y * ratio).mean() < tol:
            print("Converged after %d iters" % ii)
            break
    if ii == maxiter - 1:
        print("Did not converge !")
    Y = ratio * Y
    Z = Z / ratio
    return Y, Z
//  [python]
def newton_schultz(A, sol, maxiter=100, tol=1e-4, eps=0.1):
    """Matrix square root and inverse square root for SPD matrices"""
    d = len(A)
    normA = torch.linalg.norm(A)
    Y = A / ((1 + eps) * normA)
    Id = torch.eye(d, dtype=A.dtype, device=A.device)
    normId = torch.linalg.norm(Id)
    ratio = (normA / normId) ** 0.5

    Z = Id / ((1 + eps) * normId)
    for ii in range(maxiter):
        T = 0.5 * (3 * Id - Z @ Y)
        Y = Y @ T
        Z = T @ Z
        if abs(sol - Y * ratio).mean() < tol:
            print("Converged after %d iters" % ii)
            break
    if ii == maxiter - 1:
        print("Did not converge !")
    Y = ratio * Y
    Z = Z / ratio
    return Y, Z
//  [python]
def forward(self, x, v):
        vec1 = torch.norm(self.vec1_proj(v), dim=-2)
        vec2 = self.vec2_proj(v)

        x = torch.cat([x, vec1], dim=-1)
        x, v = torch.split(self.update_net(x), self.out_channels, dim=-1)
        v = v.unsqueeze(-2) * vec2

        if self.act is not None:
            x = self.act(x)
        return x, v
//  [python]
def forward(self, x, v):
        vec1 = torch.linalg.norm(self.vec1_proj(v), dim=-2)
        vec2 = self.vec2_proj(v)

        x = torch.cat([x, vec1], dim=-1)
        x, v = torch.split(self.update_net(x), self.out_channels, dim=-1)
        v = v.unsqueeze(-2) * vec2

        if self.act is not None:
            x = self.act(x)
        return x, v
//  [python]
def axisangle_to_R(v):
    """
    Convert an axis-angle vector to rotation matrix
    from https://github.com/ActiveVisionLab/nerfmm/blob/main/utils/lie_group_helper.py#L47

    Inputs:
        v: (3) or (B, 3)
    
    Outputs:
        R: (3, 3) or (B, 3, 3)
    """
    v_ndim = v.ndim
    if v_ndim==1:
        v = rearrange(v, 'c -> 1 c')
    zero = torch.zeros_like(v[:, :1]) # (B, 1)
    skew_v0 = torch.cat([zero, -v[:, 2:3], v[:, 1:2]], 1) # (B, 3)
    skew_v1 = torch.cat([v[:, 2:3], zero, -v[:, 0:1]], 1)
    skew_v2 = torch.cat([-v[:, 1:2], v[:, 0:1], zero], 1)
    skew_v = torch.stack([skew_v0, skew_v1, skew_v2], dim=1) # (B, 3, 3)

    norm_v = rearrange(torch.norm(v, dim=1)+1e-7, 'b -> b 1 1')
    eye = torch.eye(3, device=v.device)
    R = eye + (torch.sin(norm_v)/norm_v)*skew_v + \
        ((1-torch.cos(norm_v))/norm_v**2)*(skew_v@skew_v)
    if v_ndim==1:
        R = rearrange(R, '1 c d -> c d')
    return R
//  [python]
def axisangle_to_R(v):
    """
    Convert an axis-angle vector to rotation matrix
    from https://github.com/ActiveVisionLab/nerfmm/blob/main/utils/lie_group_helper.py#L47

    Inputs:
        v: (3) or (B, 3)
    
    Outputs:
        R: (3, 3) or (B, 3, 3)
    """
    v_ndim = v.ndim
    if v_ndim==1:
        v = rearrange(v, 'c -> 1 c')
    zero = torch.zeros_like(v[:, :1]) # (B, 1)
    skew_v0 = torch.cat([zero, -v[:, 2:3], v[:, 1:2]], 1) # (B, 3)
    skew_v1 = torch.cat([v[:, 2:3], zero, -v[:, 0:1]], 1)
    skew_v2 = torch.cat([-v[:, 1:2], v[:, 0:1], zero], 1)
    skew_v = torch.stack([skew_v0, skew_v1, skew_v2], dim=1) # (B, 3, 3)

    norm_v = rearrange(torch.linalg.norm(v, dim=1)+1e-7, 'b -> b 1 1')
    eye = torch.eye(3, device=v.device)
    R = eye + (torch.sin(norm_v)/norm_v)*skew_v + \
        ((1-torch.cos(norm_v))/norm_v**2)*(skew_v@skew_v)
    if v_ndim==1:
        R = rearrange(R, '1 c d -> c d')
    return R
//  [python]
def compute_aligned_from_right(body_right):
    xy_axis = XY_AXIS_GLOB
    x_axis = X_AXIS_GLOB

    body_right_x_proj = body_right[:,0:1] / (torch.norm(body_right[:,:2], dim=1, keepdim=True) + 1e-6)
    body_right_x_proj = torch.clamp(body_right_x_proj, min=-1.0, max=1.0) # avoid acos error

    world2aligned_angle = torch.acos(body_right_x_proj) # project to world x axis, and compute angle
    body_right = body_right * xy_axis
    world2aligned_axis = torch.cross(body_right, x_axis.expand_as(body_right))

    world2aligned_aa = (world2aligned_axis / (torch.norm(world2aligned_axis, dim=1, keepdim=True) + 1e-6)) * world2aligned_angle
    world2aligned_mat = batch_rodrigues(world2aligned_aa)

    return world2aligned_mat, world2aligned_aa
//  [python]
def compute_aligned_from_right(body_right):
    xy_axis = XY_AXIS_GLOB
    x_axis = X_AXIS_GLOB

    body_right_x_proj = body_right[:,0:1] / (torch.linalg.norm(body_right[:,:2], dim=1, keepdim=True) + 1e-6)
    body_right_x_proj = torch.clamp(body_right_x_proj, min=-1.0, max=1.0) # avoid acos error

    world2aligned_angle = torch.acos(body_right_x_proj) # project to world x axis, and compute angle
    body_right = body_right * xy_axis
    world2aligned_axis = torch.cross(body_right, x_axis.expand_as(body_right))

    world2aligned_aa = (world2aligned_axis / (torch.linalg.norm(world2aligned_axis, dim=1, keepdim=True) + 1e-6)) * world2aligned_angle
    world2aligned_mat = batch_rodrigues(world2aligned_aa)

    return world2aligned_mat, world2aligned_aa
//  [python]
def _normalize(tensor, dim=-1):
    return nan_to_num(
        torch.div(tensor, torch.norm(tensor, dim=dim, keepdim=True)))
//  [python]
def _normalize(tensor, dim=-1):
    return nan_to_num(
        torch.div(tensor, torch.linalg.norm(tensor, dim=dim, keepdim=True)))
//  [python]
def forward(self, pos, batch):
        edge_index = radius_graph(
            pos,
            r=self.cutoff_upper,
            batch=batch,
            loop=self.loop,
            max_num_neighbors=self.max_num_neighbors,
        )
        edge_vec = pos[edge_index[0]] - pos[edge_index[1]]

        if self.loop:
            # mask out self loops when computing distances because
            # the norm of 0 produces NaN gradients
            # NOTE: might influence force predictions as self loop gradients are ignored
            mask = edge_index[0] != edge_index[1]
            edge_weight = torch.zeros(edge_vec.size(0)).to(edge_vec)
            edge_weight[mask] = torch.norm(edge_vec[mask], dim=-1)
        else:
            edge_weight = torch.norm(edge_vec, dim=-1)

        lower_mask = edge_weight >= self.cutoff_lower
        edge_index = edge_index[:, lower_mask]
        edge_weight = edge_weight[lower_mask]

        if self.return_vecs:
            edge_vec = edge_vec[lower_mask]
            return edge_index, edge_weight, edge_vec
        # TODO: return only `edge_index` and `edge_weight` once
        # Union typing works with TorchScript (https://github.com/pytorch/pytorch/pull/53180)
        return edge_index, edge_weight, None
//  [python]
def forward(self, pos, batch):
        edge_index = radius_graph(
            pos,
            r=self.cutoff_upper,
            batch=batch,
            loop=self.loop,
            max_num_neighbors=self.max_num_neighbors,
        )
        edge_vec = pos[edge_index[0]] - pos[edge_index[1]]

        if self.loop:
            # mask out self loops when computing distances because
            # the norm of 0 produces NaN gradients
            # NOTE: might influence force predictions as self loop gradients are ignored
            mask = edge_index[0] != edge_index[1]
            edge_weight = torch.zeros(edge_vec.size(0)).to(edge_vec)
            edge_weight[mask] = torch.linalg.norm(edge_vec[mask], dim=-1)
        else:
            edge_weight = torch.linalg.norm(edge_vec, dim=-1)

        lower_mask = edge_weight >= self.cutoff_lower
        edge_index = edge_index[:, lower_mask]
        edge_weight = edge_weight[lower_mask]

        if self.return_vecs:
            edge_vec = edge_vec[lower_mask]
            return edge_index, edge_weight, edge_vec
        # TODO: return only `edge_index` and `edge_weight` once
        # Union typing works with TorchScript (https://github.com/pytorch/pytorch/pull/53180)
        return edge_index, edge_weight, None
//  [python]
def get_colmap_depth(
        self,
        img_p3d_all,
        img_2d_all,
        img_err_all,
        pose,
        intrinsic,
        img_w,
        img_h,
        device=0,
    ):
        # return depth and weights for each image
        # calculate normalize factor
        grid = create_meshgrid(
            img_h, img_w, normalized_coordinates=False, device=device
        )[0]
        i, j = grid.unbind(-1)
        fx, fy, cx, cy = (
            intrinsic[0, 0],
            intrinsic[1, 1],
            intrinsic[0, 2],
            intrinsic[1, 2],
        )
        directions = torch.stack(
            [(i - cx) / fx, (j - cy) / fy, torch.ones_like(i)], dim=-1
        )  # (H, W, 3)

        # Rotate ray directions from camera coordinate to the world coordinate
        rays_d = directions @ pose[:, :3].T  # (H, W, 3)
        dir_norm = torch.norm(rays_d, dim=-1, keepdim=True).reshape(img_h, img_w)

        # depth from sfm key points
        depth_all = torch.zeros(img_h, img_w, device=device)
        weights_all = torch.zeros(img_h, img_w, device=device)

        img_2d_all = torch.round(img_2d_all).long()  # (width, height)
        valid_mask = (
            (img_2d_all[:, 0] >= 0)
            & (img_2d_all[:, 0] < img_w)
            & (img_2d_all[:, 1] >= 0)
            & (img_2d_all[:, 1] < img_h)
        )

        img_2d = img_2d_all[valid_mask]
        img_err = img_err_all[valid_mask].squeeze()

        img_p3d = img_p3d_all[valid_mask]
        pose = torch.cat((pose, torch.zeros(1, 4, device=device)), dim=0)
        pose[3, 3] = 1
        extrinsic = torch.linalg.inv(pose)

        Err_mean = torch.mean(img_err)
        projected = intrinsic @ extrinsic[:3] @ img_p3d.T

        depth = projected[2, :]
        weight = 2 * torch.exp(-((img_err / Err_mean) ** 2))

        depth_all[img_2d[:, 1], img_2d[:, 0]] = depth
        weights_all[img_2d[:, 1], img_2d[:, 0]] = weight
        return depth_all.cpu() * dir_norm.cpu(), weights_all.cpu()
//  [python]
def get_colmap_depth(
        self,
        img_p3d_all,
        img_2d_all,
        img_err_all,
        pose,
        intrinsic,
        img_w,
        img_h,
        device=0,
    ):
        # return depth and weights for each image
        # calculate normalize factor
        grid = create_meshgrid(
            img_h, img_w, normalized_coordinates=False, device=device
        )[0]
        i, j = grid.unbind(-1)
        fx, fy, cx, cy = (
            intrinsic[0, 0],
            intrinsic[1, 1],
            intrinsic[0, 2],
            intrinsic[1, 2],
        )
        directions = torch.stack(
            [(i - cx) / fx, (j - cy) / fy, torch.ones_like(i)], dim=-1
        )  # (H, W, 3)

        # Rotate ray directions from camera coordinate to the world coordinate
        rays_d = directions @ pose[:, :3].T  # (H, W, 3)
        dir_norm = torch.linalg.norm(rays_d, dim=-1, keepdim=True).reshape(img_h, img_w)

        # depth from sfm key points
        depth_all = torch.zeros(img_h, img_w, device=device)
        weights_all = torch.zeros(img_h, img_w, device=device)

        img_2d_all = torch.round(img_2d_all).long()  # (width, height)
        valid_mask = (
            (img_2d_all[:, 0] >= 0)
            & (img_2d_all[:, 0] < img_w)
            & (img_2d_all[:, 1] >= 0)
            & (img_2d_all[:, 1] < img_h)
        )

        img_2d = img_2d_all[valid_mask]
        img_err = img_err_all[valid_mask].squeeze()

        img_p3d = img_p3d_all[valid_mask]
        pose = torch.cat((pose, torch.zeros(1, 4, device=device)), dim=0)
        pose[3, 3] = 1
        extrinsic = torch.linalg.inv(pose)

        Err_mean = torch.mean(img_err)
        projected = intrinsic @ extrinsic[:3] @ img_p3d.T

        depth = projected[2, :]
        weight = 2 * torch.exp(-((img_err / Err_mean) ** 2))

        depth_all[img_2d[:, 1], img_2d[:, 0]] = depth
        weights_all[img_2d[:, 1], img_2d[:, 0]] = weight
        return depth_all.cpu() * dir_norm.cpu(), weights_all.cpu()
//  [python]
def place_dihedral(
        self,
        a,
        b,
        c,
        bond_angle,
        bond_length,
        torsion_angle,
    ):
        """
        Place the point d such that the bond angle, length, and torsion angle are satisfied
        with the series a, b, c, d.
        """
        unit_vec = lambda x: x / torch.norm(x, dim=-1, keepdim=True)
        ab = b - a
        bc = unit_vec(c - b)
        d = torch.cat(
            [
                -bond_length * torch.cos(bond_angle),
                bond_length * torch.cos(torsion_angle) * torch.sin(bond_angle),
                bond_length * torch.sin(torsion_angle) * torch.sin(bond_angle),
            ],
            dim=-1
        )
        n = unit_vec(torch.cross(ab, bc))
        nbc = torch.cross(n, bc)
        m = torch.stack([bc, nbc, n], dim=1)
        d = torch.einsum("bxn,bx->bn", m, d)
        return d + c
//  [python]
def place_dihedral(
        self,
        a,
        b,
        c,
        bond_angle,
        bond_length,
        torsion_angle,
    ):
        """
        Place the point d such that the bond angle, length, and torsion angle are satisfied
        with the series a, b, c, d.
        """
        unit_vec = lambda x: x / torch.linalg.norm(x, dim=-1, keepdim=True)
        ab = b - a
        bc = unit_vec(c - b)
        d = torch.cat(
            [
                -bond_length * torch.cos(bond_angle),
                bond_length * torch.cos(torsion_angle) * torch.sin(bond_angle),
                bond_length * torch.sin(torsion_angle) * torch.sin(bond_angle),
            ],
            dim=-1
        )
        n = unit_vec(torch.cross(ab, bc))
        nbc = torch.cross(n, bc)
        m = torch.stack([bc, nbc, n], dim=1)
        d = torch.einsum("bxn,bx->bn", m, d)
        return d + c
//  [python]
def _is_valid_correlation_matrix_cholesky_factor(L, tol=1e-6):
    """Check if supplied matrix is a Cholesky factor of a valid correlation matrix

    A matrix is a Cholesky fator of a valid correlation matrix if it is lower
    triangular, has positive diagonal, and unit row-sum

    Args:
        L: A n x n lower-triangular matrix, or a batch of b lower-triangular
            matrices with shape b x n x n
        tol: The tolerance with which to check positivity of the diagonal and
            unit-sum of the rows

    Returns:
        True if L is a Cholesky factor of a valid correlation matrix, False
            otherwise (in batch mode, all matrices in the batch need to be
            Cholesky factors of valid correlation matrices)

    """
    unit_row_length = torch.all((torch.norm(L, dim=-1) - 1).abs() < tol)
    return unit_row_length and torch.all(constraints.lower_cholesky.check(L))
//  [python]
def _is_valid_correlation_matrix_cholesky_factor(L, tol=1e-6):
    """Check if supplied matrix is a Cholesky factor of a valid correlation matrix

    A matrix is a Cholesky fator of a valid correlation matrix if it is lower
    triangular, has positive diagonal, and unit row-sum

    Args:
        L: A n x n lower-triangular matrix, or a batch of b lower-triangular
            matrices with shape b x n x n
        tol: The tolerance with which to check positivity of the diagonal and
            unit-sum of the rows

    Returns:
        True if L is a Cholesky factor of a valid correlation matrix, False
            otherwise (in batch mode, all matrices in the batch need to be
            Cholesky factors of valid correlation matrices)

    """
    unit_row_length = torch.all((torch.linalg.norm(L, dim=-1) - 1).abs() < tol)
    return unit_row_length and torch.all(constraints.lower_cholesky.check(L))
//  [python]
def fit_cylinder_torch(points, normals, weights, show_warning=False):
    # compute
    # U, s, V = torch.svd(weights * normals)
    weighted_normals = weights * normals

    if np.linalg.cond(weighted_normals.data.cpu().numpy()) > 1e5:
        if show_warning:
            print("condition number is large in cylinder")
            print(torch.sum(normals).item(), torch.sum(points).item(), torch.sum(weights).item())

    U, s, V = customsvd(weighted_normals)
    a = V[:, -1]
    a = torch.reshape(a, (1, 3))

    # find the projection onto a plane perpendicular to the axis
    a = a.reshape((3, 1))
    a = a / (torch.norm(a, 2) + EPS)

    prj_circle = points - ((points @ a).permute(1, 0) * a).permute(1, 0)
    # torch doesn't have least square for

    #center, radius = fit_sphere_torch(prj_circle, normals, weights)
    center, radius = fit_circle_numpy(prj_circle.data.cpu().numpy())
    if 0:
        colors = np.array([[1, 1, 0]])  # Cylinder, yellow
        colors = np.tile(colors, [points.shape[0], 1])
        write_ply('test_cylinder.ply', points.data.numpy(), normal=normals.data.numpy(), color=colors)
        write_ply('test_cylinder_circle.ply', prj_circle.data.numpy(), normal=normals.data.numpy(), color=colors)
        write_ply('test_circle_center.ply', center1, color=colors[0:1])
    
    return a, center, radius
//  [python]
def fit_cylinder_torch(points, normals, weights, show_warning=False):
    # compute
    # U, s, V = torch.svd(weights * normals)
    weighted_normals = weights * normals

    if np.linalg.cond(weighted_normals.data.cpu().numpy()) > 1e5:
        if show_warning:
            print("condition number is large in cylinder")
            print(torch.sum(normals).item(), torch.sum(points).item(), torch.sum(weights).item())

    U, s, V = customsvd(weighted_normals)
    a = V[:, -1]
    a = torch.reshape(a, (1, 3))

    # find the projection onto a plane perpendicular to the axis
    a = a.reshape((3, 1))
    a = a / (torch.linalg.norm(a, 2) + EPS)

    prj_circle = points - ((points @ a).permute(1, 0) * a).permute(1, 0)
    # torch doesn't have least square for

    #center, radius = fit_sphere_torch(prj_circle, normals, weights)
    center, radius = fit_circle_numpy(prj_circle.data.cpu().numpy())
    if 0:
        colors = np.array([[1, 1, 0]])  # Cylinder, yellow
        colors = np.tile(colors, [points.shape[0], 1])
        write_ply('test_cylinder.ply', points.data.numpy(), normal=normals.data.numpy(), color=colors)
        write_ply('test_cylinder_circle.ply', prj_circle.data.numpy(), normal=normals.data.numpy(), color=colors)
        write_ply('test_circle_center.ply', center1, color=colors[0:1])
    
    return a, center, radius
//  [python]
def forward(self, input):
        if self.type == DeformatorType.ID:
            return input
        input = input.view([-1, self.input_dim])
        if self.type == DeformatorType.FC:
            x1 = self.fc1(input)
            x = self.act1(self.bn1(x1))
            x2 = self.fc2(x)
            x = self.act2(self.bn2(x2))
            x3 = self.fc3(x)
            x = self.act3(self.bn3(x3))
            out = self.fc4(x)
        elif self.type == DeformatorType.LINEAR:
            out = self.linear(input)
        elif self.type == DeformatorType.PROJECTIVE:
            input_norm = torch.norm(input, dim=1, keepdim=True)
            out = self.linear(input)
            out = input_norm / torch.norm(out, dim=1, keepdim=True) * out
        elif self.type == DeformatorType.ORTHO:
            mat = torch_expm((self.log_mat_half - self.log_mat_half.transpose(0, 1)).unsqueeze(0))
            out = F.linear(input, mat)
        elif self.type == DeformatorType.RANDOM:
            self.linear = self.linear
            out = F.linear(input, self.linear)
        elif self.type == DeformatorType.DEEPER_FC:
            input_norm = torch.norm(input, dim=1, keepdim=True)
            out = self.net(input)
            out = input_norm / torch.norm(out, dim=1, keepdim=True) * out
        return out
//  [python]
def forward(self, input):
        if self.type == DeformatorType.ID:
            return input
        input = input.view([-1, self.input_dim])
        if self.type == DeformatorType.FC:
            x1 = self.fc1(input)
            x = self.act1(self.bn1(x1))
            x2 = self.fc2(x)
            x = self.act2(self.bn2(x2))
            x3 = self.fc3(x)
            x = self.act3(self.bn3(x3))
            out = self.fc4(x)
        elif self.type == DeformatorType.LINEAR:
            out = self.linear(input)
        elif self.type == DeformatorType.PROJECTIVE:
            input_norm = torch.linalg.norm(input, dim=1, keepdim=True)
            out = self.linear(input)
            out = input_norm / torch.linalg.norm(out, dim=1, keepdim=True) * out
        elif self.type == DeformatorType.ORTHO:
            mat = torch_expm((self.log_mat_half - self.log_mat_half.transpose(0, 1)).unsqueeze(0))
            out = F.linear(input, mat)
        elif self.type == DeformatorType.RANDOM:
            self.linear = self.linear
            out = F.linear(input, self.linear)
        elif self.type == DeformatorType.DEEPER_FC:
            input_norm = torch.linalg.norm(input, dim=1, keepdim=True)
            out = self.net(input)
            out = input_norm / torch.linalg.norm(out, dim=1, keepdim=True) * out
        return out
//  [python]
def plot_query_embds_per_video(self, targets):
        for targets_per_video in targets:
            video_name = targets_per_video["file_names"][0].split('/')[-2]

            pred_embds = targets_per_video['embds']  # N_pred, T_prev, C
            pred_logits = targets_per_video['logits'] # N_pred, T_prev, K
            for i, (embds, logits) in enumerate(zip(pred_embds, pred_logits)):
                s, l = logits.mean(0).max(-1)
                if s < self.apply_cls_thres:
                    continue

                embds = embds / torch.norm(embds, dim=-1, keepdim=True)
                embds = embds.clamp(min=0.)
                embds = embds / torch.max(embds, dim=-1, keepdim=True)[0].clamp(min=1e-3)
                embds_np = embds.t().cpu().numpy()
                # embds_np = (embds_np * 255).astype(np.uint8)

                save_by_category = True
                if save_by_category:
                    output_dir = self.output_dir + 'query_embds/category/' + str(int(l))
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                    output_path = os.path.join(output_dir, video_name + '_' + str(i) + '.jpg')
                else:
                    output_dir = self.output_dir + 'query_embds/video/' + video_name
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                    output_path = os.path.join(output_dir, str(i) + '.jpg')

                im = plt.imshow(embds_np, cmap='jet', vmin=0, vmax=1)
                cbar = plt.colorbar(im, ticks=[0, 0.2, 0.4, 0.6, 0.8, 1])
                plt.title(str(int(l))+' + %.2f' % s)
                plt.savefig(output_path, dpi=300,bbox_inches='tight')
                plt.clf()
//  [python]
def plot_query_embds_per_video(self, targets):
        for targets_per_video in targets:
            video_name = targets_per_video["file_names"][0].split('/')[-2]

            pred_embds = targets_per_video['embds']  # N_pred, T_prev, C
            pred_logits = targets_per_video['logits'] # N_pred, T_prev, K
            for i, (embds, logits) in enumerate(zip(pred_embds, pred_logits)):
                s, l = logits.mean(0).max(-1)
                if s < self.apply_cls_thres:
                    continue

                embds = embds / torch.linalg.norm(embds, dim=-1, keepdim=True)
                embds = embds.clamp(min=0.)
                embds = embds / torch.max(embds, dim=-1, keepdim=True)[0].clamp(min=1e-3)
                embds_np = embds.t().cpu().numpy()
                # embds_np = (embds_np * 255).astype(np.uint8)

                save_by_category = True
                if save_by_category:
                    output_dir = self.output_dir + 'query_embds/category/' + str(int(l))
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                    output_path = os.path.join(output_dir, video_name + '_' + str(i) + '.jpg')
                else:
                    output_dir = self.output_dir + 'query_embds/video/' + video_name
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                    output_path = os.path.join(output_dir, str(i) + '.jpg')

                im = plt.imshow(embds_np, cmap='jet', vmin=0, vmax=1)
                cbar = plt.colorbar(im, ticks=[0, 0.2, 0.4, 0.6, 0.8, 1])
                plt.title(str(int(l))+' + %.2f' % s)
                plt.savefig(output_path, dpi=300,bbox_inches='tight')
                plt.clf()
//  [python]
def pck(predicted, target):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape
    dis = torch.norm(predicted - target, dim=len(target.shape)-1)
    #print(dis.size())
    t = torch.Tensor([0.15]).cuda()  # threshold
    out = (dis < t).float() * 1
    return out.sum()/14.0
//  [python]
def pck(predicted, target):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape
    dis = torch.linalg.norm(predicted - target, dim=len(target.shape)-1)
    #print(dis.size())
    t = torch.Tensor([0.15]).cuda()  # threshold
    out = (dis < t).float() * 1
    return out.sum()/14.0
//  [python]
def auc(predicted, target):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape
    dis = torch.norm(predicted - target, dim=len(target.shape)-1)
    outall = 0
    #print(dis.size())
    for i in range(150):
        t = torch.Tensor([float(i)/1000]).cuda()  # threshold
        out = (dis < t).float() * 1
        outall+=out.sum()/14.0
    outall = outall/150
    return outall
//  [python]
def auc(predicted, target):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape
    dis = torch.linalg.norm(predicted - target, dim=len(target.shape)-1)
    outall = 0
    #print(dis.size())
    for i in range(150):
        t = torch.Tensor([float(i)/1000]).cuda()  # threshold
        out = (dis < t).float() * 1
        outall+=out.sum()/14.0
    outall = outall/150
    return outall
//  [python]
def distance(self, P0, P1):
        # generate all line direction vectors
        n = (P1 - P0) / torch.norm(P1 - P0, dim=1).unsqueeze(1)
//  [python]
def distance(self, P0, P1):
        # generate all line direction vectors
        n = (P1 - P0) / torch.linalg.norm(P1 - P0, dim=1, keepdims=True)