jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def _orthogonalize(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, _))`,
    """
    # TODO Consider using Q = torch.orgqr(*torch.geqrf(A)) to compute the Q of the QR _much_ faster
    # and more reliably.
    # Works on FP32/64 or complex numbers (does not work for half precision)
    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:
                logging.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 scale_error(pts1, pts2, M, scale_ratio):
    """Scale error of the essential/ fundamental matrix."""

    num_pts = pts1.shape[0]

    # get homogeneous coordinates
    hom_pts1 = torch.cat((pts1, torch.ones((num_pts, 1), device=M.device)), dim=-1)
    hom_pts2 = torch.cat((pts2, torch.ones((num_pts, 1), device=M.device)), dim=-1)

    # calculate the angle between n1 and n2
    l1 = (M.transpose(-2, -1) @ (hom_pts2.transpose(-2, -1)))[:, 0:2]
    l2 = (M@(hom_pts1.transpose(-2, -1)))[:, 0:2]

    l1_norm = torch.norm(scale_ratio*l1, dim=(-1, -2))
    l2_norm = torch.norm(l2, dim=(-1, -2))

    return abs(l1_norm - l2_norm)
//  [python]
def scale_error(pts1, pts2, M, scale_ratio):
    """Scale error of the essential/ fundamental matrix."""

    num_pts = pts1.shape[0]

    # get homogeneous coordinates
    hom_pts1 = torch.cat((pts1, torch.ones((num_pts, 1), device=M.device)), dim=-1)
    hom_pts2 = torch.cat((pts2, torch.ones((num_pts, 1), device=M.device)), dim=-1)

    # calculate the angle between n1 and n2
    l1 = (M.transpose(-2, -1) @ (hom_pts2.transpose(-2, -1)))[:, 0:2]
    l2 = (M@(hom_pts1.transpose(-2, -1)))[:, 0:2]

    l1_norm = torch.linalg.norm(scale_ratio*l1, dim=(-1, -2))
    l2_norm = torch.linalg.norm(l2, dim=(-1, -2))

    return abs(l1_norm - l2_norm)
//  [python]
def orientation_error(pts1, pts2, M, ang):
    """Orientation error calculation for E or F matrix."""
    # 2D coordinates to 3D homogeneous coordinates

    num_pts = pts1.shape[0]

    # get homogeneous coordinates
    hom_pts1 = torch.cat((pts1, torch.ones((num_pts, 1), device=M.device)), dim=-1)
    hom_pts2 = torch.cat((pts2, torch.ones((num_pts, 1), device=M.device)), dim=-1)

    # calculate the ang between n1 and n2
    l1 = M.transpose(-2, -1)@hom_pts2.transpose(-2, -1)[0:2]
    l2 = M@hom_pts1.transpose(-2, -1)[0:2]
    n1 = l1[:, 0:2, :]
    n2 = l2[:, 0:2, :]

    n1_norm = 1 / torch.norm(n1, axis=0)
    n1 = torch.dot(n1, n1_norm)

    n2_norm = 1 / torch.norm(n2, axis=0)
    n2 = torch.dot(n2, n2_norm)

    alpha = torch.arccos(n1.T.dot(n2))

    ori_error = abs(alpha - ang)

    return ori_error
//  [python]
def orientation_error(pts1, pts2, M, ang):
    """Orientation error calculation for E or F matrix."""
    # 2D coordinates to 3D homogeneous coordinates

    num_pts = pts1.shape[0]

    # get homogeneous coordinates
    hom_pts1 = torch.cat((pts1, torch.ones((num_pts, 1), device=M.device)), dim=-1)
    hom_pts2 = torch.cat((pts2, torch.ones((num_pts, 1), device=M.device)), dim=-1)

    # calculate the ang between n1 and n2
    l1 = M.transpose(-2, -1)@hom_pts2.transpose(-2, -1)[0:2]
    l2 = M@hom_pts1.transpose(-2, -1)[0:2]
    n1 = l1[:, 0:2, :]
    n2 = l2[:, 0:2, :]

    n1_norm = 1 / torch.linalg.norm(n1, axis=0)
    n1 = torch.dot(n1, n1_norm)

    n2_norm = 1 / torch.linalg.norm(n2, axis=0)
    n2 = torch.dot(n2, n2_norm)

    alpha = torch.arccos(n1.T.dot(n2))

    ori_error = abs(alpha - ang)

    return ori_error
//  [python]
def vectors2aa(vec_rest, vec_final):
    # (B, 1, 1)
    vec_final_norm = torch.norm(vec_final, dim=1, keepdim=True)
    vec_rest_norm = torch.norm(vec_rest, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(vec_rest, vec_final, dim=1)
    axis_norm = torch.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(vec_rest * vec_final, dim=1, keepdim=True) / (vec_rest_norm * vec_final_norm + 1e-8)
    sin = axis_norm / (vec_rest_norm * vec_final_norm + 1e-8)

    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    return axis, sin, cos
//  [python]
def vectors2aa(vec_rest, vec_final):
    # (B, 1, 1)
    vec_final_norm = torch.linalg.norm(vec_final, dim=1, keepdim=True)
    vec_rest_norm = torch.linalg.norm(vec_rest, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(vec_rest, vec_final, dim=1)
    axis_norm = torch.linalg.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(vec_rest * vec_final, dim=1, keepdim=True) / (vec_rest_norm * vec_final_norm + 1e-8)
    sin = axis_norm / (vec_rest_norm * vec_final_norm + 1e-8)

    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    return axis, sin, cos
//  [python]
def new_decompose_E(model):
    """
      recover rotation and translation from essential matrices without SVD
      reference: Horn, Berthold KP. Recovering baseline and orientation from essential matrix[J]. J. Opt. Soc. Am, 1990, 110.
      input: essential matrix (3, 3)
      output: two possible solutions of rotation matrices, R1, R2; translation t

    """

    # assert model.shape == (3, 3)

    # Eq.18, choose the largest of the three possible pairwise cross-products
    e1, e2, e3 = model[:, 0], model[: ,1], model[: ,2]
    bs = [torch.norm(torch.cross(e1, e2)), torch.norm(torch.cross(e2, e3)), torch.norm(torch.cross(e3, e1))]
    largest = torch.argmax(torch.stack(bs))
    bb = bs[largest]

    # sqrt(1/2 trace(EE^T))
    scale_factor = torch.sqrt(0.5 * torch.trace(model @ model.transpose(0, -1)))

    if largest == 0:
      b1 = scale_factor * torch.cross(e1, e2) / torch.norm(torch.cross(e1, e2))
    elif largest == 1:
      b1 = scale_factor * torch.cross(e2, e3) / torch.norm(torch.cross(e2, e3))
    else:
      b1 = scale_factor * torch.cross(e3, e1) / torch.norm(torch.cross(e3, e1))

    # nomalization
    b1_ = b1/torch.norm(b1)

    # skew-symmetric matrix
    t0, t1, t2 = b1
    B1 = torch.tensor([
        [0, -t2, t1],
        [t2, 0, -t0],
        [-t1, t0, 0]
        ], device=b1.device)
    # the second translation and rotation
    b2 = -b1
    B2 = -B1

    # Eq.24, recover R
    # (bb)R = Cofactors(E)^T - BE
    R1 = (matrix_cofactor_tensor(model) - B1 @ model) / (b1.dot(b1))
    R2 = (matrix_cofactor_tensor(model) - B2 @ model) / (b1.dot(b1))

    return R1, R2, b1_.unsqueeze(-1)
//  [python]
def new_decompose_E(model):
    """
      recover rotation and translation from essential matrices without SVD
      reference: Horn, Berthold KP. Recovering baseline and orientation from essential matrix[J]. J. Opt. Soc. Am, 1990, 110.
      input: essential matrix (3, 3)
      output: two possible solutions of rotation matrices, R1, R2; translation t

    """

    # assert model.shape == (3, 3)

    # Eq.18, choose the largest of the three possible pairwise cross-products
    e1, e2, e3 = model[:, 0], model[: ,1], model[: ,2]
    bs = [torch.linalg.norm(torch.cross(e1, e2)), torch.linalg.norm(torch.cross(e2, e3)), torch.linalg.norm(torch.cross(e3, e1))]
    largest = torch.argmax(torch.stack(bs))
    bb = bs[largest]

    # sqrt(1/2 trace(EE^T))
    scale_factor = torch.sqrt(0.5 * torch.trace(model @ model.transpose(0, -1)))

    if largest == 0:
      b1 = scale_factor * torch.cross(e1, e2) / torch.linalg.norm(torch.cross(e1, e2))
    elif largest == 1:
      b1 = scale_factor * torch.cross(e2, e3) / torch.linalg.norm(torch.cross(e2, e3))
    else:
      b1 = scale_factor * torch.cross(e3, e1) / torch.linalg.norm(torch.cross(e3, e1))

    # nomalization
    b1_ = b1/torch.linalg.norm(b1)

    # skew-symmetric matrix
    t0, t1, t2 = b1
    B1 = torch.tensor([
        [0, -t2, t1],
        [t2, 0, -t0],
        [-t1, t0, 0]
        ], device=b1.device)
    # the second translation and rotation
    b2 = -b1
    B2 = -B1

    # Eq.24, recover R
    # (bb)R = Cofactors(E)^T - BE
    R1 = (matrix_cofactor_tensor(model) - B1 @ model) / (b1.dot(b1))
    R2 = (matrix_cofactor_tensor(model) - B2 @ model) / (b1.dot(b1))

    return R1, R2, b1_.unsqueeze(-1)
//  [python]
def _calc_lp(grads, norm_type):
    norm = 0.0
    for grad in grads:
        grad_norm = torch.norm(grad, norm_type)
        norm += grad_norm**norm_type
    return norm
//  [python]
def _calc_lp(grads, norm_type):
    norm = 0.0
    for grad in grads:
        grad_norm = torch.linalg.norm(grad, norm_type)
        norm += grad_norm**norm_type
    return norm
//  [python]
def vectors2rotmat2(vec_rest, vec_final, dtype):
    batch_size = vec_final.shape[0]
    device = vec_final.device

    # (B, 1, 1)
    vec_final_norm = torch.norm(vec_final, dim=1, keepdim=True)
    vec_rest_norm = torch.norm(vec_rest, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(vec_rest, vec_final, dim=1)
    axis_norm = torch.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(vec_rest * vec_final, dim=1, keepdim=True) / (vec_rest_norm * vec_final_norm + 1e-8)
    sin = axis_norm / (vec_rest_norm * vec_final_norm + 1e-8)

    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 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_loc = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    return rot_mat_loc, axis, cos, sin
//  [python]
def vectors2rotmat2(vec_rest, vec_final, dtype):
    batch_size = vec_final.shape[0]
    device = vec_final.device

    # (B, 1, 1)
    vec_final_norm = torch.linalg.norm(vec_final, dim=1, keepdim=True)
    vec_rest_norm = torch.linalg.norm(vec_rest, dim=1, keepdim=True)

    # (B, 3, 1)
    axis = torch.cross(vec_rest, vec_final, dim=1)
    axis_norm = torch.linalg.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(vec_rest * vec_final, dim=1, keepdim=True) / (vec_rest_norm * vec_final_norm + 1e-8)
    sin = axis_norm / (vec_rest_norm * vec_final_norm + 1e-8)

    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 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_loc = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    return rot_mat_loc, axis, cos, sin
//  [python]
def batch_rodrigues(
    rot_vecs: Tensor,
    epsilon: float = 1e-8,
) -> Tensor:
    ''' 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, dtype = rot_vecs.device, rot_vecs.dtype

    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: Tensor,
    epsilon: float = 1e-8,
) -> Tensor:
    ''' 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, dtype = rot_vecs.device, rot_vecs.dtype

    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 get_image_prior_losses(self, inputs_jit):
        # COMPUTE total variation regularization loss
        diff1 = inputs_jit[:, :, :, :-1] - inputs_jit[:, :, :, 1:]
        diff2 = inputs_jit[:, :, :-1, :] - inputs_jit[:, :, 1:, :]
        diff3 = inputs_jit[:, :, 1:, :-1] - inputs_jit[:, :, :-1, 1:]
        diff4 = inputs_jit[:, :, :-1, :-1] - inputs_jit[:, :, 1:, 1:]

        loss_var_l2 = torch.norm(diff1) + torch.norm(diff2) + torch.norm(diff3) + torch.norm(diff4)
        loss_var_l1 = (
            (diff1.abs() / 255.0).mean()
            + (diff2.abs() / 255.0).mean()
            + (diff3.abs() / 255.0).mean()
            + (diff4.abs() / 255.0).mean()
        )
        loss_var_l1 = loss_var_l1 * 255.0
        return loss_var_l1, loss_var_l2
//  [python]
def get_image_prior_losses(self, inputs_jit):
        # COMPUTE total variation regularization loss
        diff1 = inputs_jit[:, :, :, :-1] - inputs_jit[:, :, :, 1:]
        diff2 = inputs_jit[:, :, :-1, :] - inputs_jit[:, :, 1:, :]
        diff3 = inputs_jit[:, :, 1:, :-1] - inputs_jit[:, :, :-1, 1:]
        diff4 = inputs_jit[:, :, :-1, :-1] - inputs_jit[:, :, 1:, 1:]

        loss_var_l2 = torch.linalg.norm(diff1) + torch.linalg.norm(diff2) + torch.linalg.norm(diff3) + torch.linalg.norm(diff4)
        loss_var_l1 = (
            (diff1.abs() / 255.0).mean()
            + (diff2.abs() / 255.0).mean()
            + (diff3.abs() / 255.0).mean()
            + (diff4.abs() / 255.0).mean()
        )
        loss_var_l1 = loss_var_l1 * 255.0
        return loss_var_l1, loss_var_l2
//  [python]
def dist2ca(x, mask=None, eps=1e-7):
    """ Calculates distance from each point to C-alfa.
        Inputs:
        * x: (L, 14, D)
        * mask: boolean mask of (L, 14)
        Returns unit vectors and norm. 
    """
    x = x - x[:, 1].unsqueeze(1)
    norm = torch.norm(x, dim=-1, keepdim=True)
    x_norm = x / (norm+eps)
    if mask:
        return x_norm[mask], norm[mask]
    return x_norm, norm
//  [python]
def dist2ca(x, mask=None, eps=1e-7):
    """ Calculates distance from each point to C-alfa.
        Inputs:
        * x: (L, 14, D)
        * mask: boolean mask of (L, 14)
        Returns unit vectors and norm. 
    """
    x = x - x[:, 1].unsqueeze(1)
    norm = torch.linalg.norm(x, dim=-1, keepdim=True)
    x_norm = x / (norm+eps)
    if mask:
        return x_norm[mask], norm[mask]
    return x_norm, norm
//  [python]
def axis_angle_to_quaternion(axis_angle: torch.Tensor) ->torch.Tensor:
    """
    Convert rotations given as axis/angle to quaternions.

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

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

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

    Returns:
        quaternions with real part first, as tensor of shape (..., 4).
    """
    angles = torch.linalg.norm(axis_angle, ord=2, dim=-1, keepdim=True)
    half_angles = angles * 0.5
    eps = 1e-06
    small_angles = angles.abs() < eps
    sin_half_angles_over_angles = torch.empty_like(angles)
    sin_half_angles_over_angles[~small_angles] = torch.sin(half_angles[~small_angles]) / angles[~small_angles]
    sin_half_angles_over_angles[small_angles] = 0.5 - angles[small_angles] * angles[small_angles] / 48
    quaternions = torch.cat([torch.cos(half_angles), axis_angle * sin_half_angles_over_angles], dim=-1)
    return quaternions
//  [python]
def _get_vert_rad(self, vert_pos, cam_pos, orthogonal_projection, focal_length, kwargs, cloud_idx) ->torch.Tensor:
        """
        Get point radiuses.

        These can be depending on the camera position in case of a perspective
        transform.
        """
        raster_rad = self.rasterizer.raster_settings.radius
        if kwargs.get('radius_world', False):
            return raster_rad
        if isinstance(raster_rad, torch.Tensor) and raster_rad.numel() > 1 and raster_rad.ndim > 1:
            raster_rad = raster_rad[cloud_idx]
        if orthogonal_projection:
            vert_rad = torch.ones((vert_pos.shape[0],), dtype=torch.float32, device=vert_pos.device) * raster_rad
        else:
            point_dists = torch.norm(vert_pos - cam_pos, p=2, dim=1, keepdim=False)
            vert_rad = raster_rad / focal_length * point_dists
            if isinstance(self.rasterizer.cameras, PerspectiveCameras):
                pass
            else:
                vert_rad = vert_rad / 2.0
        return vert_rad
//  [python]
def _get_vert_rad(self, vert_pos, cam_pos, orthogonal_projection, focal_length, kwargs, cloud_idx) ->torch.Tensor:
        """
        Get point radiuses.

        These can be depending on the camera position in case of a perspective
        transform.
        """
        raster_rad = self.rasterizer.raster_settings.radius
        if kwargs.get('radius_world', False):
            return raster_rad
        if isinstance(raster_rad, torch.Tensor) and raster_rad.numel() > 1 and raster_rad.ndim > 1:
            raster_rad = raster_rad[cloud_idx]
        if orthogonal_projection:
            vert_rad = torch.ones((vert_pos.shape[0],), dtype=torch.float32, device=vert_pos.device) * raster_rad
        else:
            point_dists = torch.linalg.norm(vert_pos - cam_pos, ord=2, dim=1, keepdim=False)
            vert_rad = raster_rad / focal_length * point_dists
            if isinstance(self.rasterizer.cameras, PerspectiveCameras):
                pass
            else:
                vert_rad = vert_rad / 2.0
        return vert_rad
//  [python]
def batch_get_pelvis_orient(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    batch_size = rel_pose_skeleton.shape[0]
    device = rel_pose_skeleton.device

    pelvis_child = []
    for i in range(parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    spine_final_loc = rel_pose_skeleton[:, pelvis_child[-1]].clone()  # TODO: may get worse performance if choose other child ?
    spine_rest_loc = rel_rest_pose[:, pelvis_child[-1]].clone()
    spine_norm = torch.norm(spine_final_loc, dim=1, keepdim=True)
    spine_norm = spine_final_loc / (spine_norm + 1e-8)

    rot_mat_spine = vectors2rotmat(spine_rest_loc, spine_final_loc, dtype)

    assert torch.sum(torch.isnan(rot_mat_spine)
                     ) == 0, ('rot_mat_spine', rot_mat_spine, spine_final_loc, rel_pose_skeleton)
    center_final_loc = 0
    center_rest_loc = 0
    for child in pelvis_child:
        if child == pelvis_child[-1]:
            continue
        center_final_loc = center_final_loc + rel_pose_skeleton[:, child].clone()
        center_rest_loc = center_rest_loc + rel_rest_pose[:, child].clone()

    center_final_loc = center_final_loc / (len(pelvis_child) - 1)
    center_rest_loc = center_rest_loc / (len(pelvis_child) - 1)

    center_rest_loc = torch.matmul(rot_mat_spine, center_rest_loc)

    center_final_loc = center_final_loc - torch.sum(center_final_loc * spine_norm, dim=1, keepdim=True) * spine_norm
    center_rest_loc = center_rest_loc - torch.sum(center_rest_loc * spine_norm, dim=1, keepdim=True) * spine_norm

    center_final_loc_norm = torch.norm(center_final_loc, dim=1, keepdim=True)
    center_rest_loc_norm = torch.norm(center_rest_loc, dim=1, keepdim=True)

    # print(center_final_loc_norm, center_rest_loc_norm)

    # (B, 3, 1)
    axis = torch.cross(center_rest_loc, center_final_loc, dim=1)
    axis_norm = torch.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(center_rest_loc * center_final_loc, dim=1, keepdim=True) / \
        (center_rest_loc_norm * center_final_loc_norm + 1e-8)
    sin = axis_norm / (center_rest_loc_norm * center_final_loc_norm + 1e-8)

    assert torch.sum(torch.isnan(cos)
                     ) == 0, ('cos', cos)
    assert torch.sum(torch.isnan(sin)
                     ) == 0, ('sin', sin)
    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 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_center = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    rot_mat = torch.matmul(rot_mat_center, rot_mat_spine)

    return rot_mat
//  [python]
def batch_get_pelvis_orient(rel_pose_skeleton, rel_rest_pose, parents, children, dtype):
    batch_size = rel_pose_skeleton.shape[0]
    device = rel_pose_skeleton.device

    pelvis_child = []
    for i in range(parents.shape[0]):
        if parents[i] == 0 and i not in pelvis_child:
            pelvis_child.append(i)

    spine_final_loc = rel_pose_skeleton[:, pelvis_child[-1]].clone()  # TODO: may get worse performance if choose other child ?
    spine_rest_loc = rel_rest_pose[:, pelvis_child[-1]].clone()
    spine_norm = torch.linalg.norm(spine_final_loc, dim=1, keepdim=True)
    spine_norm = spine_final_loc / (spine_norm + 1e-8)

    rot_mat_spine = vectors2rotmat(spine_rest_loc, spine_final_loc, dtype)

    assert torch.sum(torch.isnan(rot_mat_spine)
                     ) == 0, ('rot_mat_spine', rot_mat_spine, spine_final_loc, rel_pose_skeleton)
    center_final_loc = 0
    center_rest_loc = 0
    for child in pelvis_child:
        if child == pelvis_child[-1]:
            continue
        center_final_loc = center_final_loc + rel_pose_skeleton[:, child].clone()
        center_rest_loc = center_rest_loc + rel_rest_pose[:, child].clone()

    center_final_loc = center_final_loc / (len(pelvis_child) - 1)
    center_rest_loc = center_rest_loc / (len(pelvis_child) - 1)

    center_rest_loc = torch.matmul(rot_mat_spine, center_rest_loc)

    center_final_loc = center_final_loc - torch.sum(center_final_loc * spine_norm, dim=1, keepdim=True) * spine_norm
    center_rest_loc = center_rest_loc - torch.sum(center_rest_loc * spine_norm, dim=1, keepdim=True) * spine_norm

    center_final_loc_norm = torch.linalg.norm(center_final_loc, dim=1, keepdim=True)
    center_rest_loc_norm = torch.linalg.norm(center_rest_loc, dim=1, keepdim=True)

    # print(center_final_loc_norm, center_rest_loc_norm)

    # (B, 3, 1)
    axis = torch.cross(center_rest_loc, center_final_loc, dim=1)
    axis_norm = torch.linalg.norm(axis, dim=1, keepdim=True)

    # (B, 1, 1)
    cos = torch.sum(center_rest_loc * center_final_loc, dim=1, keepdim=True) / \
        (center_rest_loc_norm * center_final_loc_norm + 1e-8)
    sin = axis_norm / (center_rest_loc_norm * center_final_loc_norm + 1e-8)

    assert torch.sum(torch.isnan(cos)
                     ) == 0, ('cos', cos)
    assert torch.sum(torch.isnan(sin)
                     ) == 0, ('sin', sin)
    # (B, 3, 1)
    axis = axis / (axis_norm + 1e-8)

    # Convert location revolve to rot_mat by rodrigues
    # (B, 1, 1)
    rx, ry, rz = torch.split(axis, 1, dim=1)
    zeros = torch.zeros((batch_size, 1, 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_center = ident + sin * K + (1 - cos) * torch.bmm(K, K)

    rot_mat = torch.matmul(rot_mat_center, rot_mat_spine)

    return rot_mat
//  [python]
def _koopman_matreg_loss(self):
        matreg_loss = torch.norm(self._phi_affine, p=2) + torch.norm(self._u_affine, p=2)
        return matreg_loss
//  [python]
def _koopman_matreg_loss(self):
        matreg_loss = torch.linalg.norm(self._phi_affine, p=2) + torch.linalg.norm(self._u_affine, p=2)
        return matreg_loss
//  [python]
def compute_multi_step(xt, all_xT, model, et_coeff, et_prevsum_coeff, T, t, image_dim, xT, **kwargs):
    with torch.no_grad():
        xt_all = torch.zeros_like(all_xT)
        xt_all[kwargs['xT_idx']] = xT
        xt_all[kwargs['prev_idx']] = xt[kwargs['next_idx']]

        xt = xt_all.to('cuda')

        et = model(xt, t)
        et_updated = et_coeff * et
        et_cumsum_all = et_updated.cumsum(dim=0)
        et_prevsum = et_cumsum_all

        all_seqs = torch.arange(T-1, et_cumsum_all.shape[0]-1, T)
        prev_cumsum = 0
        if len(all_seqs) > 0:
            for idx in all_seqs:
                prev_cumsum += torch.unsqueeze(et_cumsum_all[idx], dim=0)
                et_prevsum[idx+1:idx+1+T] -= torch.repeat_interleave(prev_cumsum, T,  dim=0)

        xt_next = all_xT + et_prevsum_coeff * et_prevsum
        log_dict = {
            "xt": torch.norm(xt.reshape(image_dim[0], -1), -1).mean(),
            "xt_next": torch.norm(xt_next.reshape(image_dim[0], -1), -1).mean(),
            "prediction et": torch.norm(et.reshape(image_dim[0], -1), -1).mean(),
        }
    return xt_next.view(xt.shape[0], -1), log_dict
//  [python]
def compute_multi_step(xt, all_xT, model, et_coeff, et_prevsum_coeff, T, t, image_dim, xT, **kwargs):
    with torch.no_grad():
        xt_all = torch.zeros_like(all_xT)
        xt_all[kwargs['xT_idx']] = xT
        xt_all[kwargs['prev_idx']] = xt[kwargs['next_idx']]

        xt = xt_all.to('cuda')

        et = model(xt, t)
        et_updated = et_coeff * et
        et_cumsum_all = et_updated.cumsum(dim=0)
        et_prevsum = et_cumsum_all

        all_seqs = torch.arange(T-1, et_cumsum_all.shape[0]-1, T)
        prev_cumsum = 0
        if len(all_seqs) > 0:
            for idx in all_seqs:
                prev_cumsum += torch.unsqueeze(et_cumsum_all[idx], dim=0)
                et_prevsum[idx+1:idx+1+T] -= torch.repeat_interleave(prev_cumsum, T,  dim=0)

        xt_next = all_xT + et_prevsum_coeff * et_prevsum
        log_dict = {
            "xt": torch.linalg.norm(xt.reshape(image_dim[0], -1), ord=-1).mean(),
            "xt_next": torch.linalg.norm(xt_next.reshape(image_dim[0], -1), ord=-1).mean(),
            "prediction et": torch.linalg.norm(et.reshape(image_dim[0], -1), ord=-1).mean(),
        }
    return xt_next.view(xt.shape[0], -1), log_dict
//  [python]
def get_twist(rot_mats, joints, parents):
    joints = torch.unsqueeze(joints, dim=-1)
    rel_joints = joints.clone()
    rel_joints[:, 1:] -= joints[:, parents[1:]].clone()

    # modified by xuchao
    childs = -torch.ones((parents.shape[0]), dtype=parents.dtype, device=parents.device)
    for i in range(1, parents.shape[0]):
        childs[parents[i]] = i

    dtype = rot_mats.dtype
    batch_size = rot_mats.shape[0]
    device = rot_mats.device

    angle_twist = []
    # error = False
    for i in range(1, rot_mats.shape[1]):
        # modified by xuchao
        if childs[i] < 0:
            angle_twist.append(torch.zeros((batch_size, 1), dtype=rot_mats.dtype, device=rot_mats.device))
            continue

        u = rel_joints[:, childs[i]]
        rot = rot_mats[:, i]

        v = torch.matmul(rot, u)

        u_norm = torch.norm(u, dim=1, keepdim=True)
        v_norm = torch.norm(v, dim=1, keepdim=True)

        axis = torch.cross(u, v, dim=1)
        axis_norm = torch.norm(axis, dim=1, keepdim=True)

        # (B, 1, 1)
        cos = torch.sum(u * v, dim=1, keepdim=True) / (u_norm * v_norm + 1e-8)
        sin = axis_norm / (u_norm * v_norm + 1e-8)

        # (B, 3, 1)
        axis = axis / (axis_norm + 1e-8)

        # Convert location revolve to rot_mat by rodrigues
        # (B, 1, 1)
        rx, ry, rz = torch.split(axis, 1, dim=1)
        zeros = torch.zeros((batch_size, 1, 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_pivot = ident + sin * K + (1 - cos) * torch.bmm(K, K)

        rot_mat_twist = torch.matmul(rot_mat_pivot.transpose(1, 2), rot)
        # _, axis, angle = rotmat_to_aa(rot_mat_twist)
        rot_aa_twist = matrix_to_axis_angle(rot_mat_twist)
        angle = torch.norm(rot_aa_twist, dim=1, keepdim=True)
        axis = rot_aa_twist / angle

        axis = axis / torch.norm(axis, dim=1, keepdim=True)
        spin_axis = u / u_norm
        spin_axis = spin_axis.squeeze(-1)

        pos = torch.norm(spin_axis - axis, dim=1)
        neg = torch.norm(spin_axis + axis, dim=1)

        if float(neg) < float(pos):
            if float(pos) > 1.8:
                angle_twist.append(-1 * angle)
            else:
                angle_twist.append(torch.ones_like(angle) * -999)
                print('error', float(pos), i)
        else:
            if float(neg) > 1.8:
                angle_twist.append(angle)
            else:
                angle_twist.append(torch.ones_like(angle) * -999)
                print('error', float(neg), i)

    angle_twist = torch.stack(angle_twist, dim=1)

    return angle_twist
//  [python]
def get_twist(rot_mats, joints, parents):
    joints = torch.unsqueeze(joints, dim=-1)
    rel_joints = joints.clone()
    rel_joints[:, 1:] -= joints[:, parents[1:]].clone()

    # modified by xuchao
    childs = -torch.ones((parents.shape[0]), dtype=parents.dtype, device=parents.device)
    for i in range(1, parents.shape[0]):
        childs[parents[i]] = i

    dtype = rot_mats.dtype
    batch_size = rot_mats.shape[0]
    device = rot_mats.device

    angle_twist = []
    # error = False
    for i in range(1, rot_mats.shape[1]):
        # modified by xuchao
        if childs[i] < 0:
            angle_twist.append(torch.zeros((batch_size, 1), dtype=rot_mats.dtype, device=rot_mats.device))
            continue

        u = rel_joints[:, childs[i]]
        rot = rot_mats[:, i]

        v = torch.matmul(rot, u)

        u_norm = torch.linalg.norm(u, dim=1, keepdim=True)
        v_norm = torch.linalg.norm(v, dim=1, keepdim=True)

        axis = torch.cross(u, v, dim=1)
        axis_norm = torch.linalg.norm(axis, dim=1, keepdim=True)

        # (B, 1, 1)
        cos = torch.sum(u * v, dim=1, keepdim=True) / (u_norm * v_norm + 1e-8)
        sin = axis_norm / (u_norm * v_norm + 1e-8)

        # (B, 3, 1)
        axis = axis / (axis_norm + 1e-8)

        # Convert location revolve to rot_mat by rodrigues
        # (B, 1, 1)
        rx, ry, rz = torch.split(axis, 1, dim=1)
        zeros = torch.zeros((batch_size, 1, 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_pivot = ident + sin * K + (1 - cos) * torch.bmm(K, K)

        rot_mat_twist = torch.matmul(rot_mat_pivot.transpose(1, 2), rot)
        # _, axis, angle = rotmat_to_aa(rot_mat_twist)
        rot_aa_twist = matrix_to_axis_angle(rot_mat_twist)
        angle = torch.linalg.norm(rot_aa_twist, dim=1, keepdim=True)
        axis = rot_aa_twist / angle

        axis = axis / torch.linalg.norm(axis, dim=1, keepdim=True)
        spin_axis = u / u_norm
        spin_axis = spin_axis.squeeze(-1)

        pos = torch.linalg.norm(spin_axis - axis, dim=1)
        neg = torch.linalg.norm(spin_axis + axis, dim=1)

        if float(neg) < float(pos):
            if float(pos) > 1.8:
                angle_twist.append(-1 * angle)
            else:
                angle_twist.append(torch.ones_like(angle) * -999)
                print('error', float(pos), i)
        else:
            if float(neg) > 1.8:
                angle_twist.append(angle)
            else:
                angle_twist.append(torch.ones_like(angle) * -999)
                print('error', float(neg), i)

    angle_twist = torch.stack(angle_twist, dim=1)

    return angle_twist
//  [python]
def mel_spectrogram(self, y):
        """Computes mel-spectrograms from a batch of waves
        PARAMS
        ------
        y: Variable(torch.FloatTensor) with shape (B, T) in range [-1, 1]

        RETURNS
        -------
        mel_output: torch.FloatTensor of shape (B, n_mel_channels, T)
        """
        assert torch.min(y.data) >= -1
        assert torch.max(y.data) <= 1
        magnitudes, phases = self.stft_fn.transform(y)
        magnitudes = magnitudes.data
        mel_output = torch.matmul(self.mel_basis, magnitudes)
        mel_output = self.spectral_normalize(mel_output)
        energy = torch.norm(magnitudes, dim=1)
        return mel_output, energy
//  [python]
def mel_spectrogram(self, y):
        """Computes mel-spectrograms from a batch of waves
        PARAMS
        ------
        y: Variable(torch.FloatTensor) with shape (B, T) in range [-1, 1]

        RETURNS
        -------
        mel_output: torch.FloatTensor of shape (B, n_mel_channels, T)
        """
        assert torch.min(y.data) >= -1
        assert torch.max(y.data) <= 1
        magnitudes, phases = self.stft_fn.transform(y)
        magnitudes = magnitudes.data
        mel_output = torch.matmul(self.mel_basis, magnitudes)
        mel_output = self.spectral_normalize(mel_output)
        energy = torch.linalg.norm(magnitudes, dim=1)
        return mel_output, energy
//  [python]
def orient_aa(x, mask=None, eps=1e-7):
    """ Calculates unit vectors and norms of features for backbone.
        Inputs:
        * x: (L, 14, D). Cordinates in Sidechainnet format.
        Returns unit vectors (5) and norms (3). 
    """
    # get tensor info
    device, precise = x.device, x.type()

    vec_wrap  = torch.zeros(5, x.shape[0], 3, device=device) # (feats, L, dims+1)
    norm_wrap = torch.zeros(3, x.shape[0], device=device)
    # first feat is CB-CA
    vec_wrap[0]  = x[:, 4] - x[:, 1]
    norm_wrap[0] = torch.norm(vec_wrap[0], dim=-1)
    vec_wrap[0] /= norm_wrap[0].unsqueeze(dim=-1) + eps
    # second is CA+ - CA :
    vec_wrap[1, :-1]  = x[:-1, 1] - x[1:, 1]
    norm_wrap[1, :-1] = torch.norm(vec_wrap[1, :-1], dim=-1)
    vec_wrap[1, :-1] /= norm_wrap[1, :-1].unsqueeze(dim=-1) + eps
    # same but reverse vectors
    vec_wrap[2] = (-1)*vec_wrap[1]
    # third is CA - CA-
    vec_wrap[3, 1:]  = x[:-1, 1] - x[1:, 1]
    norm_wrap[2, 1:] = torch.norm(vec_wrap[3, 1:], dim=-1)
    vec_wrap[3, 1:] /= norm_wrap[2, 1:].unsqueeze(dim=-1) + eps
    # now vectors in reverse order
    vec_wrap[4] = (-1)*vec_wrap[3]

    return vec_wrap, norm_wrap
//  [python]
def orient_aa(x, mask=None, eps=1e-7):
    """ Calculates unit vectors and norms of features for backbone.
        Inputs:
        * x: (L, 14, D). Cordinates in Sidechainnet format.
        Returns unit vectors (5) and norms (3). 
    """
    # get tensor info
    device, precise = x.device, x.type()

    vec_wrap  = torch.zeros(5, x.shape[0], 3, device=device) # (feats, L, dims+1)
    norm_wrap = torch.zeros(3, x.shape[0], device=device)
    # first feat is CB-CA
    vec_wrap[0]  = x[:, 4] - x[:, 1]
    norm_wrap[0] = torch.linalg.norm(vec_wrap[0], dim=-1)
    vec_wrap[0] /= norm_wrap[0].unsqueeze(dim=-1) + eps
    # second is CA+ - CA :
    vec_wrap[1, :-1]  = x[:-1, 1] - x[1:, 1]
    norm_wrap[1, :-1] = torch.linalg.norm(vec_wrap[1, :-1], dim=-1)
    vec_wrap[1, :-1] /= norm_wrap[1, :-1].unsqueeze(dim=-1) + eps
    # same but reverse vectors
    vec_wrap[2] = (-1)*vec_wrap[1]
    # third is CA - CA-
    vec_wrap[3, 1:]  = x[:-1, 1] - x[1:, 1]
    norm_wrap[2, 1:] = torch.linalg.norm(vec_wrap[3, 1:], dim=-1)
    vec_wrap[3, 1:] /= norm_wrap[2, 1:].unsqueeze(dim=-1) + eps
    # now vectors in reverse order
    vec_wrap[4] = (-1)*vec_wrap[3]

    return vec_wrap, norm_wrap
//  [python]
def _setweights(self):
        for name_w in self.weights:
            name_g = name_w + self.append_g
            name_v = name_w + self.append_v
            g = getattr(self.module, name_g)
            v = getattr(self.module, name_v)
            w = v*(g/torch.norm(v)).expand_as(v)
            setattr(self.module, name_w, w)
//  [python]
def _setweights(self):
        for name_w in self.weights:
            name_g = name_w + self.append_g
            name_v = name_w + self.append_v
            g = getattr(self.module, name_g)
            v = getattr(self.module, name_v)
            w = v*(g/torch.linalg.norm(v)).expand_as(v)
            setattr(self.module, name_w, w)
//  [python]
def _reset(self):
        for name_w in self.weights:
            w = getattr(self.module, name_w)

            # construct g,v such that w = g/||v|| * v
            g = torch.norm(w)
            v = w/g.expand_as(w)
            g = nn.Parameter(g.data)
            v = nn.Parameter(v.data)
            name_g = name_w + self.append_g
            name_v = name_w + self.append_v

            # remove w from parameter list
            del self.module._parameters[name_w]

            # add g and v as new parameters
            self.module.register_parameter(name_g, g)
            self.module.register_parameter(name_v, v)
//  [python]
def _reset(self):
        for name_w in self.weights:
            w = getattr(self.module, name_w)

            # construct g,v such that w = g/||v|| * v
            g = torch.linalg.norm(w)
            v = w/g.expand_as(w)
            g = nn.Parameter(g.data)
            v = nn.Parameter(v.data)
            name_g = name_w + self.append_g
            name_v = name_w + self.append_v

            # remove w from parameter list
            del self.module._parameters[name_w]

            # add g and v as new parameters
            self.module.register_parameter(name_g, g)
            self.module.register_parameter(name_v, v)
//  [python]
def forward(self, x):
        """
        Given x, returns z and the log-determinant log|df/dx|.
        Returns
        -------
        """
        if self.h in (F.elu, F.leaky_relu):
            u = self.u
        elif self.h == torch.tanh:
            scal = torch.log(1+torch.exp(self.w @ self.u)) - self.w @ self.u - 1
            u = self.u + scal * self.w / torch.norm(self.w) ** 2
        else:
            raise NotImplementedError("Non-linearity is not supported.")
        lin = torch.unsqueeze(x @ self.w, 1) + self.b
        z = x + u * self.h(lin)
        phi = functional_derivatives[self.h](lin) * self.w
        log_det = torch.log(torch.abs(1 + phi @ u) + 1e-4)
        return z, log_det
//  [python]
def forward(self, x):
        """
        Given x, returns z and the log-determinant log|df/dx|.
        Returns
        -------
        """
        if self.h in (F.elu, F.leaky_relu):
            u = self.u
        elif self.h == torch.tanh:
            scal = torch.log(1+torch.exp(self.w @ self.u)) - self.w @ self.u - 1
            u = self.u + scal * self.w / torch.linalg.norm(self.w) ** 2
        else:
            raise NotImplementedError("Non-linearity is not supported.")
        lin = torch.unsqueeze(x @ self.w, 1) + self.b
        z = x + u * self.h(lin)
        phi = functional_derivatives[self.h](lin) * self.w
        log_det = torch.log(torch.abs(1 + phi @ u) + 1e-4)
        return z, log_det
//  [python]
def _construct_model(
        device: torch.device,
        distributed_config: Optional[FSDPShampooConfig],
    ) -> Tuple[nn.Module, nn.Module, torch.Tensor, torch.Tensor]:
        data = torch.arange(16 * 4, dtype=torch.float, device=device)
        data /= torch.norm(data)
        # NOTE: We construct the model here specifically in order to ensure that
        #       FSDP Shampoo and default Shampoo produce equivalent results.
        #       This requires us to construct a model such that FSDP will split the
        #       parameters such that the resulting blocks from tensor block recovery
        #       and merging/blocking are equivalent to what would be obtained by
        #       merging/blocking on the original parameters.
        #
        #       An additional constraint imposed by FSDP is from PT2; the split must be
        #       16-byte aligned. With FP32 elements, this corresponds to 4 elements.
        #
        #       Based on the design of the model below, the model has 512 + 72 + 576 = 1160
        #       elements, which means that the model will be split at index 580 across the
        #       flattened param in FSDP.
        #       This corresponds to index 580 - 512 = 68 in the second parameter. Note that
        #       splitting at this index is equivalent to the standard blocking with a block
        #       size of 4.
        model = nn.Sequential(
            nn.Linear(16 * 4, 8, bias=False),  # 512 elements
            nn.Linear(8, 9, bias=False),  # 72 elements
            nn.Linear(9, 16 * 4, bias=False),  # 576 elements
        ).to(device=device)
        model[0].weight.data.fill_(0.1)
        model[1].weight.data.fill_(0.1)
        model[2].weight.data.fill_(0.1)
        loss = nn.MSELoss()
        target = torch.tensor([0.0]).to(device=device)
        if isinstance(distributed_config, FSDPShampooConfig):
            model = FSDP(model, use_orig_params=True)
            distributed_config.param_to_metadata = compile_fsdp_parameter_metadata(
                model
            )
            assert sum(param.numel() for param in model.parameters()) == 1160 / 2
        return model, loss, data, target
//  [python]
def _construct_model(
        device: torch.device,
        distributed_config: Optional[FSDPShampooConfig],
    ) -> Tuple[nn.Module, nn.Module, torch.Tensor, torch.Tensor]:
        data = torch.arange(16 * 4, dtype=torch.float, device=device)
        data /= torch.linalg.norm(data)
        # NOTE: We construct the model here specifically in order to ensure that
        #       FSDP Shampoo and default Shampoo produce equivalent results.
        #       This requires us to construct a model such that FSDP will split the
        #       parameters such that the resulting blocks from tensor block recovery
        #       and merging/blocking are equivalent to what would be obtained by
        #       merging/blocking on the original parameters.
        #
        #       An additional constraint imposed by FSDP is from PT2; the split must be
        #       16-byte aligned. With FP32 elements, this corresponds to 4 elements.
        #
        #       Based on the design of the model below, the model has 512 + 72 + 576 = 1160
        #       elements, which means that the model will be split at index 580 across the
        #       flattened param in FSDP.
        #       This corresponds to index 580 - 512 = 68 in the second parameter. Note that
        #       splitting at this index is equivalent to the standard blocking with a block
        #       size of 4.
        model = nn.Sequential(
            nn.Linear(16 * 4, 8, bias=False),  # 512 elements
            nn.Linear(8, 9, bias=False),  # 72 elements
            nn.Linear(9, 16 * 4, bias=False),  # 576 elements
        ).to(device=device)
        model[0].weight.data.fill_(0.1)
        model[1].weight.data.fill_(0.1)
        model[2].weight.data.fill_(0.1)
        loss = nn.MSELoss()
        target = torch.tensor([0.0]).to(device=device)
        if isinstance(distributed_config, FSDPShampooConfig):
            model = FSDP(model, use_orig_params=True)
            distributed_config.param_to_metadata = compile_fsdp_parameter_metadata(
                model
            )
            assert sum(param.numel() for param in model.parameters()) == 1160 / 2
        return model, loss, data, target
//  [python]
def norm_pointcloud(pc):
    # pc: [B, N, 3]
    center = torch.mean(pc, dim=1, keepdim=True)
    pc = pc - center
    m = torch.norm(pc, dim=-1, keepdim=True)
    pc_max = torch.max(m, dim=1, keepdim=True)[0]
    pc = pc / pc_max
    return pc
//  [python]
def norm_pointcloud(pc):
    # pc: [B, N, 3]
    center = torch.mean(pc, dim=1, keepdim=True)
    pc = pc - center
    m = torch.linalg.norm(pc, dim=-1, keepdim=True)
    pc_max = torch.max(m, dim=1, keepdim=True)[0]
    pc = pc / pc_max
    return pc
//  [python]
def get_kmeans_mu(self, x, n_centers, init_times=50, min_delta=1e-3):
        """
        Find an initial value for the mean. Requires a threshold min_delta for the k-means algorithm to stop iterating.
        The algorithm is repeated init_times often, after which the best centerpoint is returned.
        args:
            x:            torch.FloatTensor (n, d) or (n, 1, d)
            init_times:   init
            min_delta:    int
        """
        if len(x.size()) == 3:
            x = x.squeeze(1)
        x_min, x_max = x.min(), x.max()
        x = (x - x_min) / (x_max - x_min)
        
        min_cost = np.inf


        for i in range(init_times):
            tmp_center = x[np.random.choice(np.arange(x.shape[0]), size=n_centers, replace=False), ...]
            l2_dis = torch.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - tmp_center), p=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)

            cost = 0
            for c in range(n_centers):
                cost += torch.norm(x[l2_cls == c] - tmp_center[c], p=2, dim=1).mean()

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

        while delta > min_delta:
            l2_dis = torch.norm((x.unsqueeze(1).repeat(1, n_centers, 1) - center), p=2, dim=2)
            l2_cls = torch.argmin(l2_dis, dim=1)
            center_old = center.clone()

            for c in range(n_centers):
                center[c] = x[l2_cls == c].mean(dim=0)

            delta = torch.norm((center_old - center), dim=1).max()

        return (center.unsqueeze(0)*(x_max - x_min) + x_min)