jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def clip_grad_norm_(params, max_norm, aggregate_norm_fn=None) -> torch.Tensor:
    def grad_exists(p):
        return p is not None and getattr(p, "grad", None) is not None

    if isinstance(params, torch.Tensor):
        params = [params]
    params = list(params)
    grads = [
        p.grad.detach() for p in params if grad_exists(p) and not hasattr(p, "expert")
    ]
    expert_grads = [
        p.grad.detach() for p in params if grad_exists(p) and hasattr(p, "expert")
    ]

    if len(grads) == 0:
        if len(params) > 0:
            return params[0].new_tensor(0.0)
        else:
            return torch.tensor(0.0)

    if len(grads) == 1:
        total_norm = torch.norm(grads[0], p=2, dtype=torch.float32)
    else:
        if multi_tensor_l2norm_available:
            total_norm = multi_tensor_total_norm(grads)
        else:
            if torch.cuda.is_available():
                warnings.warn(
                    "amp_C fused kernels unavailable, disabling multi_tensor_l2norm; "
                    "you may get better performance by installing NVIDIA's apex library"
                )
                device = torch.cuda.current_device()
            elif grads[0].device.type == "xla":
                device = grads[0].device
            else:
                device = torch.device("cpu")
            total_norm = torch.norm(
                torch.stack(
                    [torch.norm(g, p=2, dtype=torch.float32).to(device) for g in grads]
                )
            )

    if aggregate_norm_fn is not None:
        total_norm = aggregate_norm_fn(total_norm)

    if max_norm > 0:
        max_norm = float(max_norm)
        clip_coef = (max_norm / (total_norm + 1e-6)).clamp_(max=1)
        torch._foreach_mul_(grads + expert_grads, clip_coef)

    return total_norm
//  [python]
def clip_grad_norm_(params, max_norm, aggregate_norm_fn=None) -> torch.Tensor:
    def grad_exists(p):
        return p is not None and getattr(p, "grad", None) is not None

    if isinstance(params, torch.Tensor):
        params = [params]
    params = list(params)
    grads = [
        p.grad.detach() for p in params if grad_exists(p) and not hasattr(p, "expert")
    ]
    expert_grads = [
        p.grad.detach() for p in params if grad_exists(p) and hasattr(p, "expert")
    ]

    if len(grads) == 0:
        if len(params) > 0:
            return params[0].new_tensor(0.0)
        else:
            return torch.tensor(0.0)

    if len(grads) == 1:
        total_norm = torch.linalg.norm(grads[0], ord=2, dtype=torch.float32)
    else:
        if multi_tensor_l2norm_available:
            total_norm = multi_tensor_total_norm(grads)
        else:
            if torch.cuda.is_available():
                warnings.warn(
                    "amp_C fused kernels unavailable, disabling multi_tensor_l2norm; "
                    "you may get better performance by installing NVIDIA's apex library"
                )
                device = torch.cuda.current_device()
            elif grads[0].device.type == "xla":
                device = grads[0].device
            else:
                device = torch.device("cpu")
            total_norm = torch.linalg.norm(
                torch.stack(
                    [torch.linalg.norm(g, ord=2, dtype=torch.float32).to(device) for g in grads]
                )
            )

    if aggregate_norm_fn is not None:
        total_norm = aggregate_norm_fn(total_norm)

    if max_norm > 0:
        max_norm = float(max_norm)
        clip_coef = (max_norm / (total_norm + 1e-6)).clamp_(max=1)
        torch._foreach_mul_(grads + expert_grads, clip_coef)

    return total_norm
//  [python]
def forward(ctx, M):
        normM = th.norm(M,dim=[1,2]).reshape(M.size(0),1,1)
        I = th.eye(M.size(1), requires_grad=False, device=M.device).reshape(1,M.size(1),M.size(1)).repeat(M.size(0),1,1).type(M.dtype)
        #M_sqrt = taylor_expansion(M/normM,I)
        M_sqrt_inv = pade_approximant_inverse(M / normM, I)
        M_sqrt_inv = M_sqrt_inv / th.sqrt(normM)
        ctx.save_for_backward(M, M_sqrt_inv,  I)
        return M_sqrt_inv
//  [python]
def forward(ctx, M):
        normM = th.linalg.norm(M,dim=[1,2]).reshape(M.size(0),1,1)
        I = th.eye(M.size(1), requires_grad=False, device=M.device).reshape(1,M.size(1),M.size(1)).repeat(M.size(0),1,1).type(M.dtype)
        #M_sqrt = taylor_expansion(M/normM,I)
        M_sqrt_inv = pade_approximant_inverse(M / normM, I)
        M_sqrt_inv = M_sqrt_inv / th.sqrt(normM)
        ctx.save_for_backward(M, M_sqrt_inv,  I)
        return M_sqrt_inv
//  [python]
def rot_from_axisangle(vec):
    """Convert an axisangle rotation into a 4x4 transformation matrix
    (adapted from https://github.com/Wallacoloo/printipi)
    Input 'vec' has to be Bx1x3
    """
    angle = torch.norm(vec, 2, 2, True)
    axis = vec / (angle + 1e-7)

    ca = torch.cos(angle)
    sa = torch.sin(angle)
    C = 1 - ca

    x = axis[..., 0].unsqueeze(1)
    y = axis[..., 1].unsqueeze(1)
    z = axis[..., 2].unsqueeze(1)

    xs = x * sa
    ys = y * sa
    zs = z * sa
    xC = x * C
    yC = y * C
    zC = z * C
    xyC = x * yC
    yzC = y * zC
    zxC = z * xC

    rot = torch.zeros((vec.shape[0], 4, 4)).to(device=vec.device)

    rot[:, 0, 0] = torch.squeeze(x * xC + ca)
    rot[:, 0, 1] = torch.squeeze(xyC - zs)
    rot[:, 0, 2] = torch.squeeze(zxC + ys)
    rot[:, 1, 0] = torch.squeeze(xyC + zs)
    rot[:, 1, 1] = torch.squeeze(y * yC + ca)
    rot[:, 1, 2] = torch.squeeze(yzC - xs)
    rot[:, 2, 0] = torch.squeeze(zxC - ys)
    rot[:, 2, 1] = torch.squeeze(yzC + xs)
    rot[:, 2, 2] = torch.squeeze(z * zC + ca)
    rot[:, 3, 3] = 1

    return rot
//  [python]
def rot_from_axisangle(vec):
    """Convert an axisangle rotation into a 4x4 transformation matrix
    (adapted from https://github.com/Wallacoloo/printipi)
    Input 'vec' has to be Bx1x3
    """
    angle = torch.linalg.norm(vec, 2, 2, True)
    axis = vec / (angle + 1e-7)

    ca = torch.cos(angle)
    sa = torch.sin(angle)
    C = 1 - ca

    x = axis[..., 0].unsqueeze(1)
    y = axis[..., 1].unsqueeze(1)
    z = axis[..., 2].unsqueeze(1)

    xs = x * sa
    ys = y * sa
    zs = z * sa
    xC = x * C
    yC = y * C
    zC = z * C
    xyC = x * yC
    yzC = y * zC
    zxC = z * xC

    rot = torch.zeros((vec.shape[0], 4, 4)).to(device=vec.device)

    rot[:, 0, 0] = torch.squeeze(x * xC + ca)
    rot[:, 0, 1] = torch.squeeze(xyC - zs)
    rot[:, 0, 2] = torch.squeeze(zxC + ys)
    rot[:, 1, 0] = torch.squeeze(xyC + zs)
    rot[:, 1, 1] = torch.squeeze(y * yC + ca)
    rot[:, 1, 2] = torch.squeeze(yzC - xs)
    rot[:, 2, 0] = torch.squeeze(zxC - ys)
    rot[:, 2, 1] = torch.squeeze(yzC + xs)
    rot[:, 2, 2] = torch.squeeze(z * zC + ca)
    rot[:, 3, 3] = 1

    return rot
//  [python]
def get_A_N_from_B(B):
    # Note that normals may face "inwards"
    A = torch.norm(B, p=2, dim=0)
    N = B / A.unsqueeze(0) # Normal (unit vector)
    return A, N
//  [python]
def get_A_N_from_B(B):
    # Note that normals may face "inwards"
    A = torch.linalg.norm(B, ord=2, dim=0)
    N = B / A.unsqueeze(0) # Normal (unit vector)
    return A, N
//  [python]
def backward(ctx, grad_output):
        M, M_sqrt_inv,  I = ctx.saved_tensors
        M_inv = M_sqrt_inv.bmm(M_sqrt_inv)
        grad_lya = -M_inv.bmm(grad_output).bmm(M_inv)
        norm_sqrt_inv = th.norm(M_sqrt_inv)
        a = M_sqrt_inv / norm_sqrt_inv
        q = grad_lya / norm_sqrt_inv
        for i in range(8):
            a_2 = a.bmm(a)
            q = 0.5 * (q.bmm(3.0*I-a_2)-a_2.bmm(q)+a.bmm(q).bmm(a))
            a = 0.5 * a.bmm(3.0 * I - a_2)
        grad_input = 0.5 * q
        return grad_input
//  [python]
def backward(ctx, grad_output):
        M, M_sqrt_inv,  I = ctx.saved_tensors
        M_inv = M_sqrt_inv.bmm(M_sqrt_inv)
        grad_lya = -M_inv.bmm(grad_output).bmm(M_inv)
        norm_sqrt_inv = th.linalg.norm(M_sqrt_inv)
        a = M_sqrt_inv / norm_sqrt_inv
        q = grad_lya / norm_sqrt_inv
        for i in range(8):
            a_2 = a.bmm(a)
            q = 0.5 * (q.bmm(3.0*I-a_2)-a_2.bmm(q)+a.bmm(q).bmm(a))
            a = 0.5 * a.bmm(3.0 * I - a_2)
        grad_input = 0.5 * q
        return grad_input
//  [python]
def intersection_with_unit_sphere(cam_pos, cam_rays, radius=1.0):
    '''
    get intersection with a centered sphere
    https://github.com/B1ueber2y/DIST-Renderer
    If no intersection, set first intersection to be the intersection
    with the plane tangent on the sphere orthogonal to the viewing axis
    If the camera is inside the unitsphere, set the first intersection to be the camera position
    Args:
        cam_pos (torch.FloatTensor):	 (1, *, 3) or (N, *, 3)
        cam_rays (torch.FloatTensor):	normalized viewing ray (N, *, 3)
        radius (float): sphere radius
    Returns:
        intersection0, intersection1: (N,*,3) intersections,
    '''
    if cam_pos.ndim != cam_rays.ndim:
        cam_rays = cam_rays.view(cam_rays.shape[0], -1, 3)
        cam_pos = cam_pos.view(cam_pos.shape[0], 1, 3)

    assert(cam_pos.ndim == cam_rays.ndim)

    # distance from cam_rays to center
    p, q = cam_pos, cam_rays  # (N,*,3), (N,*,3)
    ptq = (p * q).sum(dim=-1)  # (N,*)
    # middle point between the two intersections
    mid = p - ptq[..., None] * q  # (N,*,3)
    dist = torch.norm(mid, p=2, dim=-1)  # (N, *)
    cam_dist = torch.norm(p, dim=-1)
    valid_mask = (dist <= radius)
    value = radius ** 2 - dist ** 2

    # length between the intersections, valid_mask indicates whether there's an intersection
    # set ~valid_mask to be intersection with the plane tangent to the unit sphere and orthogonal
    # to the camera axis
    maxbound_marching_zdepth = torch.full_like(dist, 10.0)
    maxbound_marching_zdepth[valid_mask] = 2 * torch.sqrt(value[valid_mask])

    cam_pos_dist = torch.norm(cam_pos, dim=-1)  # (N,*)

    # If the cameras are inside the sphere, then set the initial zdepth to 0
    # if torch.sum(cam_pos_dist > radius) == 0:
    #     init_zdepth = torch.zeros_like(dist)
    # else:
    # First intersection depth
    init_zdepth = torch.zeros_like(dist)
    init_zdepth_valid = torch.sqrt(
        cam_pos_dist.expand_as(dist)[valid_mask] ** 2 - dist[valid_mask] ** 2) - maxbound_marching_zdepth[valid_mask] / 2.0  # (N)
    init_zdepth[valid_mask] = init_zdepth_valid
    # If no intersection, set first intersection to be the intersection
    # with the plane tangent on the sphere orthogonal to the viewing axis
    init_zdepth[~valid_mask] = ((cam_dist - radius) /
                                eps_denom(-ptq / cam_dist))[~valid_mask]

    intersection0 = init_zdepth.unsqueeze(-1) * cam_rays + cam_pos
    intersection1 = maxbound_marching_zdepth[...,
                                             None] * cam_rays + intersection0
    invalid_far_zdepth = ((radius + cam_dist) /
                          eps_denom(-ptq / cam_dist))[~valid_mask]
    cam_pos = cam_pos.expand_as(cam_rays)
    intersection1[~valid_mask] = invalid_far_zdepth.unsqueeze(
        -1) * cam_rays[~valid_mask] + cam_pos[~valid_mask]

    return intersection0, intersection1, valid_mask
//  [python]
def intersection_with_unit_sphere(cam_pos, cam_rays, radius=1.0):
    '''
    get intersection with a centered sphere
    https://github.com/B1ueber2y/DIST-Renderer
    If no intersection, set first intersection to be the intersection
    with the plane tangent on the sphere orthogonal to the viewing axis
    If the camera is inside the unitsphere, set the first intersection to be the camera position
    Args:
        cam_pos (torch.FloatTensor):	 (1, *, 3) or (N, *, 3)
        cam_rays (torch.FloatTensor):	normalized viewing ray (N, *, 3)
        radius (float): sphere radius
    Returns:
        intersection0, intersection1: (N,*,3) intersections,
    '''
    if cam_pos.ndim != cam_rays.ndim:
        cam_rays = cam_rays.view(cam_rays.shape[0], -1, 3)
        cam_pos = cam_pos.view(cam_pos.shape[0], 1, 3)

    assert(cam_pos.ndim == cam_rays.ndim)

    # distance from cam_rays to center
    p, q = cam_pos, cam_rays  # (N,*,3), (N,*,3)
    ptq = (p * q).sum(dim=-1)  # (N,*)
    # middle point between the two intersections
    mid = p - ptq[..., None] * q  # (N,*,3)
    dist = torch.linalg.norm(mid, ord=2, dim=-1)  # (N, *)
    cam_dist = torch.linalg.norm(p, ord=2, dim=-1)
    valid_mask = (dist <= radius)
    value = radius ** 2 - dist ** 2

    # length between the intersections, valid_mask indicates whether there's an intersection
    # set ~valid_mask to be intersection with the plane tangent to the unit sphere and orthogonal
    # to the camera axis
    maxbound_marching_zdepth = torch.full_like(dist, 10.0)
    maxbound_marching_zdepth[valid_mask] = 2 * torch.sqrt(value[valid_mask])

    cam_pos_dist = torch.linalg.norm(cam_pos, ord=2, dim=-1)  # (N,*)

    # If the cameras are inside the sphere, then set the initial zdepth to 0
    # if torch.sum(cam_pos_dist > radius) == 0:
    #     init_zdepth = torch.zeros_like(dist)
    # else:
    # First intersection depth
    init_zdepth = torch.zeros_like(dist)
    init_zdepth_valid = torch.sqrt(
        cam_pos_dist.expand_as(dist)[valid_mask] ** 2 - dist[valid_mask] ** 2) - maxbound_marching_zdepth[valid_mask] / 2.0  # (N)
    init_zdepth[valid_mask] = init_zdepth_valid
    # If no intersection, set first intersection to be the intersection
    # with the plane tangent on the sphere orthogonal to the viewing axis
    init_zdepth[~valid_mask] = ((cam_dist - radius) /
                                eps_denom(-ptq / cam_dist))[~valid_mask]

    intersection0 = init_zdepth.unsqueeze(-1) * cam_rays + cam_pos
    intersection1 = maxbound_marching_zdepth[...,
                                             None] * cam_rays + intersection0
    invalid_far_zdepth = ((radius + cam_dist) /
                          eps_denom(-ptq / cam_dist))[~valid_mask]
    cam_pos = cam_pos.expand_as(cam_rays)
    intersection1[~valid_mask] = invalid_far_zdepth.unsqueeze(
        -1) * cam_rays[~valid_mask] + cam_pos[~valid_mask]

    return intersection0, intersection1, valid_mask
//  [python]
def forward(self, inputs):
        dvecs = []
        for wav in inputs:
            mel = self.get_mel(wav)
            if mel.dim() == 3:
                mel = mel.squeeze(0)
            dvecs += [self.embedder(mel)]
        dvecs = torch.stack(dvecs)

        dvec = torch.mean(dvecs, dim=0)
        dvec = dvec / torch.norm(dvec)

        return dvec
//  [python]
def forward(self, inputs):
        dvecs = []
        for wav in inputs:
            mel = self.get_mel(wav)
            if mel.dim() == 3:
                mel = mel.squeeze(0)
            dvecs += [self.embedder(mel)]
        dvecs = torch.stack(dvecs)

        dvec = torch.mean(dvecs, dim=0)
        dvec = dvec / torch.linalg.norm(dvec)

        return dvec
//  [python]
def get_mel(self, y):
        if self.pad and y.shape[-1] < 14000:
            y = F.pad(y, (0, 14000 - y.shape[-1]))

        window = torch.hann_window(self.win_length).to(y)
        y = torch.stft(y, n_fft=self.embedder_params["n_fft"],
                       hop_length=self.hop_length,
                       win_length=self.win_length,
                       window=window)
        magnitudes = torch.norm(y, dim=-1, p=2) ** 2
        mel = torch.log10(self.mel_basis @ magnitudes + 1e-6)
        return mel
//  [python]
def get_mel(self, y):
        if self.pad and y.shape[-1] < 14000:
            y = F.pad(y, (0, 14000 - y.shape[-1]))

        window = torch.hann_window(self.win_length).to(y)
        y = torch.stft(y, n_fft=self.embedder_params["n_fft"],
                       hop_length=self.hop_length,
                       win_length=self.win_length,
                       window=window)
        magnitudes = torch.linalg.norm(y, dim=-1, ord=2) ** 2
        mel = torch.log10(self.mel_basis @ magnitudes + 1e-6)
        return mel
//  [python]
def rotvec_to_rotmat(rot_vecs):
    """Convert axis-angle representation to rotation matrix.
    Args:
        rot_vecs: size = [B, 3]
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    l1norm = torch.norm(rot_vecs + 1e-8, p=2, dim=1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(rot_vecs, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * normalized], dim=1)
    return quat_to_rotmat(quat)
//  [python]
def rotvec_to_rotmat(rot_vecs):
    """Convert axis-angle representation to rotation matrix.
    Args:
        rot_vecs: size = [B, 3]
    Returns:
        Rotation matrix corresponding to the quaternion -- size = [B, 3, 3]
    """
    l1norm = torch.linalg.norm(rot_vecs + 1e-8, ord=2, dim=1)
    angle = torch.unsqueeze(l1norm, -1)
    normalized = torch.div(rot_vecs, angle)
    angle = angle * 0.5
    v_cos = torch.cos(angle)
    v_sin = torch.sin(angle)
    quat = torch.cat([v_cos, v_sin * normalized], dim=1)
    return quat_to_rotmat(quat)
//  [python]
def forward(self, mel):
        # (num_mels, T) -> (num_mels, T', window)
        mels = mel.unfold(1, self.hp["window"], self.hp["stride"])
        mels = mels.permute(1, 2, 0)  # (T', window, num_mels)
        x, _ = self.lstm(mels)  # (T', window, lstm_hidden)
        x = x[:, -1, :]  # (T', lstm_hidden), use last frame only
        x = self.proj(x)  # (T', emb_dim)
        x = x / torch.norm(x, p=2, dim=1, keepdim=True)  # (T', emb_dim)

        x = x.mean(dim=0)
        if x.norm(p=2) != 0:
            x = x / x.norm(p=2)
        return x
//  [python]
def forward(self, mel):
        # (num_mels, T) -> (num_mels, T', window)
        mels = mel.unfold(1, self.hp["window"], self.hp["stride"])
        mels = mels.permute(1, 2, 0)  # (T', window, num_mels)
        x, _ = self.lstm(mels)  # (T', window, lstm_hidden)
        x = x[:, -1, :]  # (T', lstm_hidden), use last frame only
        x = self.proj(x)  # (T', emb_dim)
        x = x / torch.linalg.norm(x, ord=2, dim=1, keepdim=True)  # (T', emb_dim)

        x = x.mean(dim=0)
        if torch.linalg.norm(x, ord=2) != 0:
            x = x / torch.linalg.norm(x, ord=2)
        return x
//  [python]
def get_slope(self, bottom_corners):
        """
        Args:
            bottom_corners: shape(N, M, 4, 2)
        Returns:
            pass
        """
        # import ipdb
        # ipdb.set_trace()
        start_corners = bottom_corners
        end_corners = bottom_corners[:, :, [1, 2, 3, 0]]
        direction = end_corners - start_corners
        norm = torch.norm(direction, dim=-1).unsqueeze(-1)
        return direction / (norm + 1e-7)
//  [python]
def get_slope(self, bottom_corners):
        """
        Args:
            bottom_corners: shape(N, M, 4, 2)
        Returns:
            pass
        """
        # import ipdb
        # ipdb.set_trace()
        start_corners = bottom_corners
        end_corners = bottom_corners[:, :, [1, 2, 3, 0]]
        direction = end_corners - start_corners
        norm = torch.linalg.norm(direction, dim=-1).unsqueeze(-1)
        return direction / (norm + 1e-7)
//  [python]
def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
    ''' Calculates the rotation matrices for a batch of rotation vectors
        Parameters
        ----------
        rot_vecs: torch.tensor Nx3
            array of N axis-angle vectors
        Returns
        -------
        R: torch.tensor Nx3x3
            The rotation matrices for the given axis-angle parameters
    '''

    batch_size = rot_vecs.shape[0]
    device = rot_vecs.device

    angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
    rot_dir = rot_vecs / angle

    cos = torch.unsqueeze(torch.cos(angle), dim=1)
    sin = torch.unsqueeze(torch.sin(angle), dim=1)

    # Bx1 arrays
    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
    K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)

    zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))

    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
    return rot_mat
//  [python]
def batch_rodrigues(rot_vecs, epsilon=1e-8, dtype=torch.float32):
    ''' Calculates the rotation matrices for a batch of rotation vectors
        Parameters
        ----------
        rot_vecs: torch.tensor Nx3
            array of N axis-angle vectors
        Returns
        -------
        R: torch.tensor Nx3x3
            The rotation matrices for the given axis-angle parameters
    '''

    batch_size = rot_vecs.shape[0]
    device = rot_vecs.device

    angle = torch.linalg.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
    rot_dir = rot_vecs / angle

    cos = torch.unsqueeze(torch.cos(angle), dim=1)
    sin = torch.unsqueeze(torch.sin(angle), dim=1)

    # Bx1 arrays
    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
    K = torch.zeros((batch_size, 3, 3), dtype=dtype, device=device)

    zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
        .view((batch_size, 3, 3))

    ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
    rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
    return rot_mat
//  [python]
def _onenorm(A):
    return torch.norm(A, 1).item()
//  [python]
def _onenorm(A):
    return torch.linalg.norm(A, 1).item()
//  [python]
def compute_predictive_dist(args, G, loader):

    R = 1000
    eps = torch.randn(R, args.epsilon_dim)
    sampled_weights = G(eps)
    wholex = loader.dataset[:][0]
    wholey = loader.dataset[:][1]
    list_of_param_dicts = weights_to_dict(args, sampled_weights)

    pred_logprob = 0
    for param_dict in list_of_param_dicts:

        mean = torch.matmul(torch.matmul(wholex, param_dict['a']), param_dict['b'])
        pred_logprob += -(args.output_dim*np.log(2 * np.pi) + torch.norm(wholey-mean,dim=1)**2) / 2

    return pred_logprob/R
//  [python]
def compute_predictive_dist(args, G, loader):

    R = 1000
    eps = torch.randn(R, args.epsilon_dim)
    sampled_weights = G(eps)
    wholex = loader.dataset[:][0]
    wholey = loader.dataset[:][1]
    list_of_param_dicts = weights_to_dict(args, sampled_weights)

    pred_logprob = 0
    for param_dict in list_of_param_dicts:

        mean = torch.matmul(torch.matmul(wholex, param_dict['a']), param_dict['b'])
        pred_logprob += -(args.output_dim*np.log(2 * np.pi) + torch.linalg.norm(wholey-mean,dim=1)**2) / 2

    return pred_logprob/R
//  [python]
def compute_Bg(pred_logprob,loader,args):

    wholex = loader.dataset[:][0]
    wholey = loader.dataset[:][1]
    mean = torch.matmul(torch.matmul(wholex, args.a_params), args.b_params)
    mean_dev = torch.norm(wholey-mean,dim=1)**2
    true_logprob = -(args.output_dim * np.log(2 * np.pi) + mean_dev) / 2

    return (true_logprob-pred_logprob).mean()
//  [python]
def compute_Bg(pred_logprob,loader,args):

    wholex = loader.dataset[:][0]
    wholey = loader.dataset[:][1]
    mean = torch.matmul(torch.matmul(wholex, args.a_params), args.b_params)
    mean_dev = torch.linalg.norm(wholey-mean,dim=1)**2
    true_logprob = -(args.output_dim * np.log(2 * np.pi) + mean_dev) / 2

    return (true_logprob-pred_logprob).mean()
//  [python]
def test_lobpcg_torchscript(self, device, dtype):
        from torch.testing._internal.common_utils import random_sparse_pd_matrix
        from torch._linalg_utils import matmul as mm

        lobpcg = torch.jit.script(torch.lobpcg)

        m = 500
        k = 5
        A1 = random_sparse_pd_matrix(m, density=2.0 / m, device=device, dtype=dtype)
        X1 = torch.randn((m, k), dtype=dtype, device=device)
        E1, V1 = lobpcg(A1, X=X1)
        eq_err = torch.norm((mm(A1, V1) - V1 * E1), 2) / E1.max()
        self.assertLess(eq_err, 1e-6)
//  [python]
def test_lobpcg_torchscript(self, device, dtype):
        from torch.testing._internal.common_utils import random_sparse_pd_matrix
        from torch._linalg_utils import matmul as mm

        lobpcg = torch.jit.script(torch.lobpcg)

        m = 500
        k = 5
        A1 = random_sparse_pd_matrix(m, density=2.0 / m, device=device, dtype=dtype)
        X1 = torch.randn((m, k), dtype=dtype, device=device)
        E1, V1 = lobpcg(A1, X=X1)
        eq_err = torch.linalg.norm((mm(A1, V1) - V1 * E1), 2) / E1.max()
        self.assertLess(eq_err, 1e-6)
//  [python]
def _orthogonalize_gram_schmidt(matrix, epsilon=0):
    """
    Applies Gram-Schmidt procedure to orthogonalize a given 2D tensor.
    If epsilon is 0, this is equivalent to `torch.qr(matrix, out=(matrix, _))`,
    """
    num_cols = matrix.shape[1]
    for i in range(num_cols):
        # Normalize the i'th column.
        col = matrix[:, i: i + 1]
        # If no epsilon is added here, division by zero may be caused by vanishing gradients.
        # This epsilon is not needed if the input matrix covers the gradients of at least one entire layer in the neural network.
        if epsilon == 0:
            # Note that col ** 2 can underflow/overflow if we use FP16.
            # May need to consider multiplying a scaling factor and dividing it later, or using bfloat16 instead.
            try:
                col /= torch.norm(col)
            except ZeroDivisionError:
                logger.error(
                    "The matrix to be orthogonalized has at least a column of all 0s. Please set a small value such as 1e-8 "
                    "as `orthogonalization_epsilon` in PowerSGD state."
                )
                # Recover the values from NaNs to 0s.
                col.fill_(0.0)
        else:
            col /= torch.norm(col) + epsilon
        # Project it on the rest and remove it.
        if i + 1 < num_cols:
            rest = matrix[:, i + 1:]
            rest -= torch.sum(col * rest, dim=0) * col
//  [python]
def _orthogonalize_gram_schmidt(matrix, epsilon=0):
    """
    Applies Gram-Schmidt procedure to orthogonalize a given 2D tensor.
    If epsilon is 0, this is equivalent to `torch.qr(matrix, out=(matrix, _))`,
    """
    num_cols = matrix.shape[1]
    for i in range(num_cols):
        # Normalize the i'th column.
        col = matrix[:, i: i + 1]
        # If no epsilon is added here, division by zero may be caused by vanishing gradients.
        # This epsilon is not needed if the input matrix covers the gradients of at least one entire layer in the neural network.
        if epsilon == 0:
            # Note that col ** 2 can underflow/overflow if we use FP16.
            # May need to consider multiplying a scaling factor and dividing it later, or using bfloat16 instead.
            try:
                col /= torch.linalg.norm(col)
            except ZeroDivisionError:
                logger.error(
                    "The matrix to be orthogonalized has at least a column of all 0s. Please set a small value such as 1e-8 "
                    "as `orthogonalization_epsilon` in PowerSGD state."
                )
                # Recover the values from NaNs to 0s.
                col.fill_(0.0)
        else:
            col /= torch.linalg.norm(col) + epsilon
        # Project it on the rest and remove it.
        if i + 1 < num_cols:
            rest = matrix[:, i + 1:]
            rest -= torch.sum(col * rest, dim=0) * col
//  [python]
def _brute_cdist(self, x, y, p=2):
        r1 = x.shape[-2]
        r2 = y.shape[-2]
        if r1 == 0 or r2 == 0:
            return torch.empty(r1, r2, device=x.device)
        return torch.norm(x[..., None, :] - y[..., None, :, :], p=p, dim=-1)
//  [python]
def _brute_cdist(self, x, y, p=2):
        r1 = x.shape[-2]
        r2 = y.shape[-2]
        if r1 == 0 or r2 == 0:
            return torch.empty(r1, r2, device=x.device)
        return torch.linalg.norm(x[..., None, :] - y[..., None, :, :], ord=p, dim=-1)
//  [python]
def test_norm_fro_2_equivalence_old(self, device, dtype):
        input_sizes = [
            (0,),
            (10,),
            (0, 0),
            (4, 30),
            (0, 45),
            (100, 0),
            (45, 10, 23),
            (0, 23, 59),
            (23, 0, 37),
            (34, 58, 0),
            (0, 0, 348),
            (0, 3434, 0),
            (0, 0, 0),
            (5, 3, 8, 1, 3, 5)]

        for input_size in input_sizes:
            a = make_tensor(input_size, device, dtype, low=-9, high=9)

            # Try full reduction
            dim_settings = [None]

            # Try all possible 1-D reductions
            dim_settings += list(range(-a.dim(), a.dim()))

            def wrap_dim(dim, ndims):
                assert (dim < ndims) and (dim >= -ndims)
                if dim >= 0:
                    return dim
                else:
                    return dim + ndims

            # Try all possible 2-D reductions
            dim_settings += [
                (d0, d1) for d0, d1 in itertools.combinations(range(-a.dim(), a.dim()), 2)
                if wrap_dim(d0, a.dim()) != wrap_dim(d1, a.dim())]

            for dim in dim_settings:
                for keepdim in [True, False]:
                    a_norm_2 = torch.norm(a, p=2, dim=dim, keepdim=keepdim)
                    a_norm_fro = torch.norm(a, p='fro', dim=dim, keepdim=keepdim)
                    self.assertEqual(a_norm_fro, a_norm_2)
//  [python]
def test_norm_fro_2_equivalence_old(self, device, dtype):
        input_sizes = [
            (0,),
            (10,),
            (0, 0),
            (4, 30),
            (0, 45),
            (100, 0),
            (45, 10, 23),
            (0, 23, 59),
            (23, 0, 37),
            (34, 58, 0),
            (0, 0, 348),
            (0, 3434, 0),
            (0, 0, 0),
            (5, 3, 8, 1, 3, 5)]

        for input_size in input_sizes:
            a = make_tensor(input_size, device, dtype, low=-9, high=9)

            # Try full reduction
            dim_settings = [None]

            # Try all possible 1-D reductions
            dim_settings += list(range(-a.dim(), a.dim()))

            def wrap_dim(dim, ndims):
                assert (dim < ndims) and (dim >= -ndims)
                if dim >= 0:
                    return dim
                else:
                    return dim + ndims

            # Try all possible 2-D reductions
            dim_settings += [
                (d0, d1) for d0, d1 in itertools.combinations(range(-a.dim(), a.dim()), 2)
                if wrap_dim(d0, a.dim()) != wrap_dim(d1, a.dim())]

            for dim in dim_settings:
                for keepdim in [True, False]:
                    a_norm_2 = torch.linalg.norm(a, ord=2, dim=dim, keepdim=keepdim)
                    a_norm_fro = torch.linalg.norm(a, ord='fro', dim=dim, keepdim=keepdim)
                    self.assertEqual(a_norm_fro, a_norm_2)
//  [python]
def approximnate_isometry(self):
        # use dictionary to apply masks.
        print('approximate_isometry starts...')
        approxed_paras = {}
        for module in self.modules:
            for name, weight in module.named_parameters():
                if name not in self.masks: continue
                approxed_paras[name] = weight
        optimizer = optim.SGD(approxed_paras.values(), momentum=0.9, lr=0.1)
        for e in range(1, 10001):
            if e == int(10000/3):
                for param_group in optimizer.param_groups:
                    param_group['lr'] = 0.01
            if e == int(10000*2/3):
                for param_group in optimizer.param_groups:
                    param_group['lr'] = 0.001
            self.apply_mask()
            optimizer.zero_grad()
            layer_loss = {}
            for name in approxed_paras:
                para = approxed_paras[name]
                layer_loss[name] = torch.norm(torch.matmul(torch.transpose(para.view(-1, para.size()[1]),0,1), para.view(-1, para.size()[1])) - torch.eye(para.size()[1]).cuda())
            loss = sum(layer_loss.values())
            # loss.to(self.device)
            loss.backward(retain_graph=True)
            optimizer.step()
            if e % 1000 == 0:
                print(e, loss.item())
        self.apply_mask()
//  [python]
def approximnate_isometry(self):
        # use dictionary to apply masks.
        print('approximate_isometry starts...')
        approxed_paras = {}
        for module in self.modules:
            for name, weight in module.named_parameters():
                if name not in self.masks: continue
                approxed_paras[name] = weight
        optimizer = optim.SGD(approxed_paras.values(), momentum=0.9, lr=0.1)
        for e in range(1, 10001):
            if e == int(10000/3):
                for param_group in optimizer.param_groups:
                    param_group['lr'] = 0.01
            if e == int(10000*2/3):
                for param_group in optimizer.param_groups:
                    param_group['lr'] = 0.001
            self.apply_mask()
            optimizer.zero_grad()
            layer_loss = {}
            for name in approxed_paras:
                para = approxed_paras[name]
                layer_loss[name] = torch.linalg.norm(torch.matmul(torch.transpose(para.view(-1, para.size()[1]),0,1), para.view(-1, para.size()[1])) - torch.eye(para.size()[1]).cuda())
            loss = sum(layer_loss.values())
            # loss.to(self.device)
            loss.backward(retain_graph=True)
            optimizer.step()
            if e % 1000 == 0:
                print(e, loss.item())
        self.apply_mask()
//  [python]
def weighted_mpjpe(predicted, target, w):
    """
    Weighted mean per-joint position error (i.e. mean Euclidean distance)
    """
    assert predicted.shape == target.shape
    # assert w.shape[0] == predicted.shape[0]
    return torch.mean(w * torch.norm(predicted - target, dim=len(target.shape)-1))
//  [python]
def weighted_mpjpe(predicted, target, w):
    """
    Weighted mean per-joint position error (i.e. mean Euclidean distance)
    """
    assert predicted.shape == target.shape
    # assert w.shape[0] == predicted.shape[0]
    return torch.mean(w * torch.linalg.norm(predicted - target, dim=len(target.shape)-1))
//  [python]
def mean_velocity_error_train(predicted, target, axis=0):
    """
    Mean per-joint velocity error (i.e. mean Euclidean distance of the 1st derivative)
    """
    assert predicted.shape == target.shape
    
    velocity_predicted = torch.diff(predicted, dim=axis)
    velocity_target = torch.diff(target, dim=axis)

    return torch.mean(torch.norm(velocity_predicted - velocity_target, dim=len(target.shape)-1))
//  [python]
def mean_velocity_error_train(predicted, target, axis=0):
    """
    Mean per-joint velocity error (i.e. mean Euclidean distance of the 1st derivative)
    """
    assert predicted.shape == target.shape
    
    velocity_predicted = torch.diff(predicted, dim=axis)
    velocity_target = torch.diff(target, dim=axis)

    return torch.mean(torch.linalg.norm(velocity_predicted - velocity_target, dim=len(target.shape)-1))
//  [python]
def mpjpe(predicted, target, return_joints_err=False):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape
    if not return_joints_err:
        return torch.mean(torch.norm(predicted - target, dim=len(target.shape)-1))
    else:
        errors = torch.norm(predicted - target, dim=len(target.shape)-1)
        # errors: [B, T, N]
        from einops import rearrange
        errors = rearrange(errors, 'B T N -> N (B T)')
        errors = torch.mean(errors, dim=-1).cpu().numpy().reshape(-1) * 1000
        return torch.mean(torch.norm(predicted - target, dim=len(target.shape)-1)), errors
//  [python]
def mpjpe(predicted, target, return_joints_err=False):
    """
    Mean per-joint position error (i.e. mean Euclidean distance),
    often referred to as "Protocol #1" in many papers.
    """
    assert predicted.shape == target.shape
    if not return_joints_err:
        return torch.mean(torch.linalg.norm(predicted - target, dim=len(target.shape)-1))
    else:
        errors = torch.linalg.norm(predicted - target, dim=len(target.shape)-1)
        # errors: [B, T, N]
        from einops import rearrange
        errors = rearrange(errors, 'B T N -> N (B T)')
        errors = torch.mean(errors, dim=-1).cpu().numpy().reshape(-1) * 1000
        return torch.mean(torch.linalg.norm(predicted - target, dim=len(target.shape)-1)), errors
//  [python]
def l1_authogonal(self, mu=0.01, density=0.5):
        """
        a python implementation solving the following problem:
        Given initial weight W, we minimize the optimization problem by optimizing M:
        min_{M}    || （M.*W)^T (M.*W) - Ir ||_1 + mu || M.*W ||_1
        where M is the target mask, which should be binary and M.*W is nearly orthogonal
        and M.*W is sparse
        """
        print('pruning at initializatin with l1norm authogonal')
        approxed_paras = copy.deepcopy(self.masks)
        for name in approxed_paras:
            approxed_paras[name].requires_grad = True
        optimizer = optim.SGD(approxed_paras.values(), momentum=0.9, lr=0.1)

        for e in range(0, 10000):
            if e == int(10000/3):
                for param_group in optimizer.param_groups:
                    param_group['lr'] = 0.01
            if e == int(10000*2/3):
                for param_group in optimizer.param_groups:
                    param_group['lr'] = 0.001

            self.apply_mask()
            optimizer.zero_grad()
            layer_loss = {}
            for name, weight in self.dense_model.named_parameters():
                if name in self.masks:
                    para = approxed_paras[name] * weight
                    layer_loss[name] = torch.norm(torch.matmul(torch.transpose(para.view(-1, para.size()[1]),0,1), para.view(-1, para.size()[1])) - torch.eye(para.size()[1]).cuda()) \
                                       + mu*torch.norm(para, 1)

            loss = sum(layer_loss.values())
            # loss.to(self.device)
            loss.backward(retain_graph=True)
            optimizer.step()
            if e % 1000 == 0:
                print(e, loss.item())


        scores = []
        for name, weight in self.dense_model.named_parameters():
            if name in self.masks:
                scores.append(torch.abs(approxed_paras[name]* weight))

        all_scores = torch.cat([torch.flatten(x) for x in scores])

        num_params_to_keep = int(len(all_scores) * density)
        threshold, _ = torch.topk(all_scores, num_params_to_keep, sorted=True)
        acceptable_score = threshold[-1]

        layer_wise_sparsities = []
        for g in scores:
            mask = (g >= acceptable_score).float()
            layer_wise_sparsities.append(mask)

        return layer_wise_sparsities
//  [python]
def l1_authogonal(self, mu=0.01, density=0.5):
        """
        a python implementation solving the following problem:
        Given initial weight W, we minimize the optimization problem by optimizing M:
        min_{M}    || （M.*W)^T (M.*W) - Ir ||_1 + mu || M.*W ||_1
        where M is the target mask, which should be binary and M.*W is nearly orthogonal
        and M.*W is sparse
        """
        print('pruning at initializatin with l1norm authogonal')
        approxed_paras = copy.deepcopy(self.masks)
        for name in approxed_paras:
            approxed_paras[name].requires_grad = True
        optimizer = optim.SGD(approxed_paras.values(), momentum=0.9, lr=0.1)

        for e in range(0, 10000):
            if e == int(10000/3):
                for param_group in optimizer.param_groups:
                    param_group['lr'] = 0.01
            if e == int(10000*2/3):
                for param_group in optimizer.param_groups:
                    param_group['lr'] = 0.001

            self.apply_mask()
            optimizer.zero_grad()
            layer_loss = {}
            for name, weight in self.dense_model.named_parameters():
                if name in self.masks:
                    para = approxed_paras[name] * weight
                    layer_loss[name] = torch.linalg.norm(torch.matmul(torch.transpose(para.view(-1, para.size()[1]),0,1), para.view(-1, para.size()[1])) - torch.eye(para.size()[1]).cuda(), ord=1) \
                                       + mu*torch.linalg.norm(para, ord=1)

            loss = sum(layer_loss.values())
            # loss.to(self.device)
            loss.backward(retain_graph=True)
            optimizer.step()
            if e % 1000 == 0:
                print(e, loss.item())


        scores = []
        for name, weight in self.dense_model.named_parameters():
            if name in self.masks:
                scores.append(torch.abs(approxed_paras[name]* weight))

        all_scores = torch.cat([torch.flatten(x) for x in scores])

        num_params_to_keep = int(len(all_scores) * density)
        threshold, _ = torch.topk(all_scores, num_params_to_keep, sorted=True)
        acceptable_score = threshold[-1]

        layer_wise_sparsities = []
        for g in scores:
            mask = (g >= acceptable_score).float()
            layer_wise_sparsities.append(mask)

        return layer_wise_sparsities
//  [python]
def forward(self, node_attr, pos, batch):
        # edge_index = radius_graph(pos, self.cutoff, batch=batch, loop=False)
        edge_index = knn_graph(pos, k=self.k, batch=batch, flow='target_to_source')
        edge_length = torch.norm(pos[edge_index[0]] - pos[edge_index[1]], dim=1)
        edge_attr = self.distance_expansion(edge_length)

        h = node_attr
        for interaction in self.interactions:
            h = h + interaction(h, edge_index, edge_attr)
        return h
//  [python]
def forward(self, node_attr, pos, batch):
        edge_index = knn_graph(pos, k=self.k, batch=batch, flow='target_to_source')
        edge_length = torch.linalg.norm(pos[edge_index[0]] - pos[edge_index[1]], dim=1)
        edge_attr = self.distance_expansion(edge_length)

        h = node_attr
        for interaction in self.interactions:
            h = h + interaction(h, edge_index, edge_attr)
        return h
//  [python]
def test_norm(self):
        a = torch.arange(9, dtype=torch.float, device="mps") - 4
        b = a.reshape((3, 3))

        a_cpu = torch.arange(9, dtype=torch.float, device="cpu") - 4
        b_cpu = a_cpu.reshape((3, 3))

        res = torch.norm(a)
        res_cpu = torch.norm(a_cpu)
        self.assertEqual(res, res_cpu)

        res = torch.norm(b)
        res_cpu = torch.norm(b_cpu)
        self.assertEqual(res, res_cpu)

        res = torch.norm(a, float('inf'))
        res_cpu = torch.norm(a_cpu, float('inf'))
        self.assertEqual(res, res_cpu)

        res = torch.norm(b, float('inf'))
        res_cpu = torch.norm(b_cpu, float('inf'))
        self.assertEqual(res, res_cpu)

        c = torch.tensor([[1, 2, 3], [-1, 1, 4]], dtype=torch.float, device="mps")
        c_cpu = torch.tensor([[1, 2, 3], [-1, 1, 4]] , dtype=torch.float, device="cpu")

        res = torch.norm(c, dim=0)
        res_cpu = torch.norm(c_cpu, dim=0)
        self.assertEqual(res, res_cpu)

        res = torch.norm(c, dim=1)
        res_cpu = torch.norm(c_cpu, dim=1)
        self.assertEqual(res, res_cpu)

        res = torch.norm(c, p=1, dim=1)
        res_cpu = torch.norm(c_cpu, p=1, dim=1)
        self.assertEqual(res, res_cpu)

        d = torch.arange(8, dtype=torch.float, device="mps").reshape(2, 2, 2)
        d_cpu = torch.arange(8, dtype=torch.float, device="cpu").reshape(2, 2, 2)

        res = torch.norm(d, dim=(1, 2))
        res_cpu = torch.norm(d_cpu, dim=(1, 2))
        self.assertEqual(res, res_cpu)

        res = torch.norm(d[0, :, :]), torch.norm(d[1, :, :])
        res_cpu = torch.norm(d_cpu[0, :, :]), torch.norm(d_cpu[1, :, :])
        self.assertEqual(res, res_cpu)
//  [python]
def test_norm(self):
        a = torch.arange(9, dtype=torch.float, device="mps") - 4
        b = a.reshape((3, 3))

        a_cpu = torch.arange(9, dtype=torch.float, device="cpu") - 4
        b_cpu = a_cpu.reshape((3, 3))

        res = torch.linalg.norm(a)
        res_cpu = torch.linalg.norm(a_cpu)
        self.assertEqual(res, res_cpu)

        res = torch.linalg.norm(b)
        res_cpu = torch.linalg.norm(b_cpu)
        self.assertEqual(res, res_cpu)

        res = torch.linalg.norm(a, float('inf'))
        res_cpu = torch.linalg.norm(a_cpu, float('inf'))
        self.assertEqual(res, res_cpu)

        res = torch.linalg.norm(b, float('inf'))
        res_cpu = torch.linalg.norm(b_cpu, float('inf'))
        self.assertEqual(res, res_cpu)

        c = torch.tensor([[1, 2, 3], [-1, 1, 4]], dtype=torch.float, device="mps")
        c_cpu = torch.tensor([[1, 2, 3], [-1, 1, 4]] , dtype=torch.float, device="cpu")

        res = torch.linalg.norm(c, dim=0)
        res_cpu = torch.linalg.norm(c_cpu, dim=0)
        self.assertEqual(res, res_cpu)

        res = torch.linalg.norm(c, dim=1)
        res_cpu = torch.linalg.norm(c_cpu, dim=1)
        self.assertEqual(res, res_cpu)

        res = torch.linalg.norm(c, ord=1, dim=1)
        res_cpu = torch.linalg.norm(c_cpu, ord=1, dim=1)
        self.assertEqual(res, res_cpu)

        d = torch.arange(8, dtype=torch.float, device="mps").reshape(2, 2, 2)
        d_cpu = torch.arange(8, dtype=torch.float, device="cpu").reshape(2, 2, 2)

        res = torch.linalg.norm(d, dim=(1, 2))
        res_cpu = torch.linalg.norm(d_cpu, dim=(1, 2))
        self.assertEqual(res, res_cpu)

        res = torch.linalg.norm(d[0, :, :]), torch.linalg.norm(d[1, :, :])
        res_cpu = torch.linalg.norm(d_cpu[0, :, :]), torch.linalg.norm(d_cpu[1, :, :])
        self.assertEqual(res, res_cpu)
//  [python]
def bonelen_consistency_loss(dataset,keypoints,pred_out):
    loss_length = 0
    if dataset == 'h36m':
        if keypoints.startswith('hr'):
            assert "hrnet has not completed"
        else:
            bones = [(0,1), (0,4), (1,2), (2,3), (4,5), (5,6), (0,7), (7,8), (8,9), (9,10), 
                    (8,11), (11,12), (12,13), (8,14), (14,15), (15,16)]
        for (i,j) in bones:
            bonelen = pred_out[:,:,i]-pred_out[:,:,j]
            bone_diff = bonelen[:,1:,:]-bonelen[:,:-1,:]
            loss_length += torch.mean(torch.norm(bone_diff, dim=-1))
    elif dataset.startswith('heva'):
        loss_length = 0

    return 0.01 * loss_length
//  [python]
def bonelen_consistency_loss(dataset,keypoints,pred_out):
    loss_length = 0
    if dataset == 'h36m':
        if keypoints.startswith('hr'):
            assert "hrnet has not completed"
        else:
            bones = [(0,1), (0,4), (1,2), (2,3), (4,5), (5,6), (0,7), (7,8), (8,9), (9,10), 
                    (8,11), (11,12), (12,13), (8,14), (14,15), (15,16)]
        for (i,j) in bones:
            bonelen = pred_out[:,:,i]-pred_out[:,:,j]
            bone_diff = bonelen[:,1:,:]-bonelen[:,:-1,:]
            loss_length += torch.mean(torch.linalg.norm(bone_diff, dim=-1))
    elif dataset.startswith('heva'):
        loss_length = 0

    return 0.01 * loss_length
//  [python]
def sym_penalty(dataset,keypoints,pred_out):
    """
    get penalty for the symmetry of human body
    :return:
    """
    loss_sym = 0
    if dataset == 'h36m':
        if keypoints.startswith('hr'):
            left_bone = [(0,4),(4,5),(5,6),(8,10),(10,11),(11,12)]
            right_bone = [(0,1),(1,2),(2,3),(8,13),(13,14),(14,15)]
        else:
            left_bone = [(0,4),(4,5),(5,6),(8,11),(11,12),(12,13)]
            right_bone = [(0,1),(1,2),(2,3),(8,14),(14,15),(15,16)]
        for (i_left,j_left),(i_right,j_right) in zip(left_bone,right_bone):
            left_part = pred_out[:,:,i_left]-pred_out[:,:,j_left]
            right_part = pred_out[:, :, i_right] - pred_out[:, :, j_right]
            loss_sym += torch.mean(torch.abs(torch.norm(left_part, dim=-1) - torch.norm(right_part, dim=-1)))
    elif dataset.startswith('STB'):
        loss_sym = 0
    return 0.01*loss_sym