jiahuimbzuai/CoQuIR huggingface dataset:
//  [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
//  [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.linalg.norm(
                (x.unsqueeze(1).repeat(1, n_centers, 1) - tmp_center), ord=2, dim=2
            )
            l2_cls = torch.argmin(l2_dis, dim=1)

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

            if cost < min_cost:
                min_cost = cost
                center = tmp_center

        delta = np.inf

        while delta > min_delta:
            l2_dis = torch.linalg.norm(
                (x.unsqueeze(1).repeat(1, n_centers, 1) - center), ord=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.linalg.norm((center_old - center), ord=2, dim=1).max()

        return center.unsqueeze(0) * (x_max - x_min) + x_min
//  [python]
def _patch_loss(self, evaluator: Evaluator):
        def loss_patch(original_loss, batch):
            reg_loss = 0.
            count = 0
            for module_name, target_scores in self.scores.items():
                for target_name, score in target_scores.items():
                    target_space = self._target_spaces[module_name][target_name]
                    if target_space.sparse_threshold is not None:
                        reg_loss += torch.norm(score.sigmoid(), p=1) / score.numel()  # type: ignore
                        count += 1
            ratio = max(0., min(1., 1 - (self._remaining_times / self.total_times) ** 3))
            if count > 0:
                reg_loss = self.regular_scale * ratio * reg_loss / count
            return original_loss + reg_loss

        evaluator.patch_loss(loss_patch)
//  [python]
def _patch_loss(self, evaluator: Evaluator):
        def loss_patch(original_loss, batch):
            reg_loss = 0.
            count = 0
            for module_name, target_scores in self.scores.items():
                for target_name, score in target_scores.items():
                    target_space = self._target_spaces[module_name][target_name]
                    if target_space.sparse_threshold is not None:
                        reg_loss += torch.linalg.norm(score.sigmoid(), ord=1) / score.numel()  # type: ignore
                        count += 1
            ratio = max(0., min(1., 1 - (self._remaining_times / self.total_times) ** 3))
            if count > 0:
                reg_loss = self.regular_scale * ratio * reg_loss / count
            return original_loss + reg_loss

        evaluator.patch_loss(loss_patch)
//  [python]
def correlation(X):
    """
    Weird definition of correlation used in Cohen, Chung, Lee and Sampolinsky (2020)
    Divides by the vector norm instead of the standard deviation
    """
    X_c = X - X.mean(dim=1).unsqueeze(-1)
    stds = torch.norm(X_c, dim=1)
    
    # X_c = X
    # stds = torch.norm(X, dim=1)

    corr_mat = torch.einsum('ij,kj->ik', X_c, X_c)/torch.einsum('i,j->ij', stds, stds)

    return corr_mat
//  [python]
def correlation(X):
    """
    Weird definition of correlation used in Cohen, Chung, Lee and Sampolinsky (2020)
    Divides by the vector norm instead of the standard deviation
    """
    X_c = X - X.mean(dim=1).unsqueeze(-1)
    stds = torch.linalg.norm(X_c, dim=1)
    
    # X_c = X
    # stds = torch.linalg.norm(X, dim=1)

    corr_mat = torch.einsum('ij,kj->ik', X_c, X_c)/torch.einsum('i,j->ij', stds, stds)

    return corr_mat
//  [python]
def forward(self, x):
        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
        return x / norm.clamp(min=self.eps) * self.g
//  [python]
def forward(self, x):
        norm = torch.linalg.norm(x, dim=-1, keepdim=True) * self.scale
        return x / norm.clamp(min=self.eps) * self.g
//  [python]
def forward(self, x, norm_scale_shift_inp):
        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
        norm = x / norm.clamp(min=self.eps) * self.g

        ss_emb = self.scale_shift_process(norm_scale_shift_inp)
        scale, shift = torch.chunk(ss_emb, 2, dim=1)
        h = norm * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
        return h
//  [python]
def forward(self, x, norm_scale_shift_inp):
        norm = torch.linalg.norm(x, dim=-1, keepdim=True) * self.scale
        norm = x / norm.clamp(min=self.eps) * self.g

        ss_emb = self.scale_shift_process(norm_scale_shift_inp)
        scale, shift = torch.chunk(ss_emb, 2, dim=1)
        h = norm * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
        return h
//  [python]
def create_cost_and_derivates(manifold, matrix, backend):
    euclidean_gradient = euclidean_hessian = None

    if backend == "autograd":

        @pymanopt.function.autograd(manifold)
        def cost(X):
            return 0.25 * np.linalg.norm(X.T @ X - matrix) ** 2

    elif backend == "jax":

        @pymanopt.function.jax(manifold)
        def cost(X):
            return 0.25 * jnp.linalg.norm(X.T @ X - matrix) ** 2

    elif backend == "numpy":

        @pymanopt.function.numpy(manifold)
        def cost(X):
            return 0.25 * np.linalg.norm(X.T @ X - matrix) ** 2

        @pymanopt.function.numpy(manifold)
        def euclidean_gradient(X):
            return 0.5 * X @ (X.T @ X - matrix)

        @pymanopt.function.numpy(manifold)
        def euclidean_hessian(X, H):
            return X @ (H.T @ X + X.T @ H) + H @ (X.T @ X - matrix)

    elif backend == "pytorch":
        matrix_ = torch.from_numpy(matrix)

        @pymanopt.function.pytorch(manifold)
        def cost(X):
            return (
                0.25 * torch.norm(torch.transpose(X, 1, 0) @ X - matrix_) ** 2
            )

    elif backend == "tensorflow":

        @pymanopt.function.tensorflow(manifold)
        def cost(X):
            return 0.25 * tf.norm(tf.transpose(X) @ X - matrix) ** 2

    else:
        raise ValueError(f"Unsupported backend '{backend}'")

    return cost, euclidean_gradient, euclidean_hessian
//  [python]
def create_cost_and_derivates(manifold, matrix, backend):
    euclidean_gradient = euclidean_hessian = None

    if backend == "autograd":

        @pymanopt.function.autograd(manifold)
        def cost(X):
            return 0.25 * np.linalg.norm(X.T @ X - matrix) ** 2

    elif backend == "jax":

        @pymanopt.function.jax(manifold)
        def cost(X):
            return 0.25 * jnp.linalg.norm(X.T @ X - matrix) ** 2

    elif backend == "numpy":

        @pymanopt.function.numpy(manifold)
        def cost(X):
            return 0.25 * np.linalg.norm(X.T @ X - matrix) ** 2

        @pymanopt.function.numpy(manifold)
        def euclidean_gradient(X):
            return 0.5 * X @ (X.T @ X - matrix)

        @pymanopt.function.numpy(manifold)
        def euclidean_hessian(X, H):
            return X @ (H.T @ X + X.T @ H) + H @ (X.T @ X - matrix)

    elif backend == "pytorch":
        matrix_ = torch.from_numpy(matrix)

        @pymanopt.function.pytorch(manifold)
        def cost(X):
            return (
                0.25 * torch.linalg.norm(torch.transpose(X, 1, 0) @ X - matrix_) ** 2
            )

    elif backend == "tensorflow":

        @pymanopt.function.tensorflow(manifold)
        def cost(X):
            return 0.25 * tf.norm(tf.transpose(X) @ X - matrix) ** 2

    else:
        raise ValueError(f"Unsupported backend '{backend}'")

    return cost, euclidean_gradient, euclidean_hessian
//  [python]
def test_bro():
    batch = torch.tensor([0, 0, 0, 0, 1, 1, 1, 2, 2])

    g1 = torch.tensor([
        [0.2, 0.2, 0.2, 0.2],
        [0.0, 0.2, 0.2, 0.2],
        [0.2, 0.0, 0.2, 0.2],
        [0.2, 0.2, 0.0, 0.2],
    ])

    g2 = torch.tensor([
        [0.2, 0.2, 0.2, 0.2],
        [0.0, 0.2, 0.2, 0.2],
        [0.2, 0.0, 0.2, 0.2],
    ])

    g3 = torch.tensor([
        [0.2, 0.2, 0.2, 0.2],
        [0.2, 0.0, 0.2, 0.2],
    ])

    s = 0.
    for g in [g1, g2, g3]:
        s += torch.norm(g @ g.t() - torch.eye(g.shape[0]), p=2)

    assert torch.isclose(s / 3., bro(torch.cat([g1, g2, g3], dim=0), batch))
//  [python]
def test_bro():
    batch = torch.tensor([0, 0, 0, 0, 1, 1, 1, 2, 2])

    g1 = torch.tensor([
        [0.2, 0.2, 0.2, 0.2],
        [0.0, 0.2, 0.2, 0.2],
        [0.2, 0.0, 0.2, 0.2],
        [0.2, 0.2, 0.0, 0.2],
    ])

    g2 = torch.tensor([
        [0.2, 0.2, 0.2, 0.2],
        [0.0, 0.2, 0.2, 0.2],
        [0.2, 0.0, 0.2, 0.2],
    ])

    g3 = torch.tensor([
        [0.2, 0.2, 0.2, 0.2],
        [0.2, 0.0, 0.2, 0.2],
    ])

    s = 0.
    for g in [g1, g2, g3]:
        s += torch.linalg.norm(g @ g.t() - torch.eye(g.shape[0]), ord=2)

    assert torch.isclose(s / 3., bro(torch.cat([g1, g2, g3], dim=0), batch))
//  [python]
def norm(self, p="fro", dim=None, keepdim=False):
        r"""See :func: `torch.norm`"""
        return torch.norm(self, p, dim, keepdim)
//  [python]
def norm(self, p="fro", dim=None, keepdim=False):
        r"""See :func: `torch.norm`"""
        return torch.linalg.norm(self, p, dim, keepdim)
//  [python]
def compute_curvature_torsion(x):
    # Move input data to the GPU if available
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    x = x.to(device)

    # Compute central point and neighbor point tensors
    x_central = x[:, :, 0, :]  # Central point
    x_neighbors = x[:, :, 1:, :]  # Neighbor points

    # Compute neighbor vectors
    delta = x_neighbors - x_central.unsqueeze(2)  # Shape: (batch_size, num_points, near_feature - 1, dim)

    # Compute pairwise dot products
    dot_product = torch.sum(delta * delta, dim=-1)  # Shape: (batch_size, num_points, near_feature - 1)

    # Compute weights using softmax function
    weights = F.softmax(-dot_product, dim=2)

    # Compute curvature
    curvature = torch.sum(weights * dot_product, dim=2)  # Shape: (batch_size, num_points)

    # Compute torsion
    delta_diff = delta[:, :, :-1, :] - delta[:, :, 1:, :]
    torsion = torch.sum(weights[:, :, :-1] * torch.norm(delta_diff, dim=-1) ** 2, dim=2)  # Shape: (batch_size, num_points)

    return curvature, torsion
//  [python]
def compute_curvature_torsion(x):
    # Move input data to the GPU if available
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    x = x.to(device)

    # Compute central point and neighbor point tensors
    x_central = x[:, :, 0, :]  # Central point
    x_neighbors = x[:, :, 1:, :]  # Neighbor points

    # Compute neighbor vectors
    delta = x_neighbors - x_central.unsqueeze(2)  # Shape: (batch_size, num_points, near_feature - 1, dim)

    # Compute pairwise dot products
    dot_product = torch.sum(delta * delta, dim=-1)  # Shape: (batch_size, num_points, near_feature - 1)

    # Compute weights using softmax function
    weights = F.softmax(-dot_product, dim=2)

    # Compute curvature
    curvature = torch.sum(weights * dot_product, dim=2)  # Shape: (batch_size, num_points)

    # Compute torsion
    delta_diff = delta[:, :, :-1, :] - delta[:, :, 1:, :]
    torsion = torch.sum(weights[:, :, :-1] * torch.linalg.norm(delta_diff, dim=-1) ** 2, dim=2)  # Shape: (batch_size, num_points)

    return curvature, torsion
//  [python]
def compute_manifold_measures(X, y, diagnostic_mode=False):
    """
    Function to calculate simple measures of manifold radius and dimension
    Also returns the correlation matrix of the features belonging to different classes
    :param X: (torch.Tensor) Feature matrix (batch should be along the first dimension)
    :param y: (torch.Tensor) Label vector
    :return: (three torch.Tensors)
            class correlation matrix, manifold radii for each class, participation ratio (linear measure of
            manifold dimensionality) for each class
    """
    class_centers = []
    class_radii = []
    class_dims = []

    global_center = X.mean(dim=0)
    global_radius = torch.norm(X - global_center, dim=1).mean()

    for i in torch.unique(y):
        x = X[y == i]

        pr = participation_ratio(x)
        class_center = x.mean(dim=0)
        class_radius = torch.norm(x - class_center, dim=1).mean()

        class_centers.append(class_center)
        class_radii.append(class_radius)
        class_dims.append(pr)

    class_centers = torch.stack(class_centers)
    class_radii = torch.stack(class_radii)
    class_dims = torch.stack(class_dims)
    centroid_spreads = torch.norm(class_centers - global_center, dim=1)

    class_correlation_matrix = correlation(class_centers)

    if not diagnostic_mode:
        class_correlation = class_correlation_matrix[np.triu_indices(class_correlation_matrix.shape[0], k=1)].mean()
        return class_correlation, class_radii.mean()/global_radius, class_dims.mean(), global_radius
    else:
        return class_correlation_matrix, class_radii, class_dims, centroid_spreads, global_radius
//  [python]
def compute_manifold_measures(X, y, diagnostic_mode=False):
    """
    Function to calculate simple measures of manifold radius and dimension
    Also returns the correlation matrix of the features belonging to different classes
    :param X: (torch.Tensor) Feature matrix (batch should be along the first dimension)
    :param y: (torch.Tensor) Label vector
    :return: (three torch.Tensors)
            class correlation matrix, manifold radii for each class, participation ratio (linear measure of
            manifold dimensionality) for each class
    """
    class_centers = []
    class_radii = []
    class_dims = []

    global_center = X.mean(dim=0)
    global_radius = torch.linalg.norm(X - global_center, dim=1).mean()

    for i in torch.unique(y):
        x = X[y == i]

        pr = participation_ratio(x)
        class_center = x.mean(dim=0)
        class_radius = torch.linalg.norm(x - class_center, dim=1).mean()

        class_centers.append(class_center)
        class_radii.append(class_radius)
        class_dims.append(pr)

    class_centers = torch.stack(class_centers)
    class_radii = torch.stack(class_radii)
    class_dims = torch.stack(class_dims)
    centroid_spreads = torch.linalg.norm(class_centers - global_center, dim=1)

    class_correlation_matrix = correlation(class_centers)

    if not diagnostic_mode:
        class_correlation = class_correlation_matrix[np.triu_indices(class_correlation_matrix.shape[0], k=1)].mean()
        return class_correlation, class_radii.mean()/global_radius, class_dims.mean(), global_radius
    else:
        return class_correlation_matrix, class_radii, class_dims, centroid_spreads, global_radius
//  [python]
def _normal_vector(coordinates: torch.Tensor, angle_idxs: torch.Tensor) -> torch.Tensor:
    vector1 = coordinates[angle_idxs[:, 0]] - coordinates[angle_idxs[:, 1]]
    vector2 = coordinates[angle_idxs[:, 2]] - coordinates[angle_idxs[:, 1]]

    cross_product = torch.cross(vector1, vector2)

    return cross_product / torch.norm(cross_product, dim=-1).unsqueeze(-1)
//  [python]
def _normal_vector(coordinates: torch.Tensor, angle_idxs: torch.Tensor) -> torch.Tensor:
    vector1 = coordinates[angle_idxs[:, 0]] - coordinates[angle_idxs[:, 1]]
    vector2 = coordinates[angle_idxs[:, 2]] - coordinates[angle_idxs[:, 1]]

    cross_product = torch.cross(vector1, vector2)

    return cross_product / torch.linalg.norm(cross_product, dim=-1, keepdim=True)
//  [python]
def calc_drmsd(chain_a, chain_b, use_gpu=False):
    assert len(chain_a) == len(chain_b)
    distance_matrix_a = calc_pairwise_distances(chain_a, chain_a, use_gpu)
    distance_matrix_b = calc_pairwise_distances(chain_b, chain_b, use_gpu)
    return torch.norm(distance_matrix_a - distance_matrix_b, 2) \
           / math.sqrt((len(chain_a) * (len(chain_a) - 1)))
//  [python]
def calc_drmsd(chain_a, chain_b, use_gpu=False):
    assert len(chain_a) == len(chain_b)
    distance_matrix_a = calc_pairwise_distances(chain_a, chain_a, use_gpu)
    distance_matrix_b = calc_pairwise_distances(chain_b, chain_b, use_gpu)
    return torch.linalg.norm(distance_matrix_a - distance_matrix_b, 2) \
           / math.sqrt((len(chain_a) * (len(chain_a) - 1)))
//  [python]
def _in_sphere(x, r, eps):
    norm = x.norm(dim=-1)
    rs = torch.full_like(norm, r)
    return (torch.norm(norm - rs, p=float('inf')) < eps).all()
//  [python]
def _in_sphere(x, r, eps):
    norm = torch.linalg.norm(x, dim=-1)
    rs = torch.full_like(norm, r)
    return (torch.linalg.norm(norm - rs, ord=float('inf')) < eps).all()
//  [python]
def quaternion_to_axis_angle(quaternions: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as quaternions to axis/angle.

    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).

    Returns:
        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.
    """
    norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
    half_angles = torch.atan2(norms, quaternions[..., :1])
    angles = 2 * half_angles
    eps = 1e-6
    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]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    return quaternions[..., 1:] / sin_half_angles_over_angles
//  [python]
def quaternion_to_axis_angle(quaternions: torch.Tensor) -> torch.Tensor:
    """
    Convert rotations given as quaternions to axis/angle.

    Args:
        quaternions: quaternions with real part first,
            as tensor of shape (..., 4).

    Returns:
        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.
    """
    norms = torch.linalg.norm(quaternions[..., 1:], ord=2, dim=-1, keepdim=True)
    half_angles = torch.atan2(norms, quaternions[..., :1])
    angles = 2 * half_angles
    eps = 1e-6
    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]
    )
    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
    # so sin(x/2)/x is about 1/2 - (x*x)/48
    sin_half_angles_over_angles[small_angles] = (
        0.5 - (angles[small_angles] * angles[small_angles]) / 48
    )
    return quaternions[..., 1:] / sin_half_angles_over_angles
//  [python]
def remove_back_points(self, frame):
        rel_pose = frame.get_rel_pose()
        points = frame.get_points()
        points_norm = torch.norm(points, 2, -1)
        points_xy = points[:, :2]
        if rel_pose == None:
            x = 1
            y = 0
        else:
            x = rel_pose[0, 3]
            y = rel_pose[1, 3]
        rel_xy = torch.ones((1, 2))
        rel_xy[0, 0] = x
        rel_xy[0, 1] = y
        point_cos = torch.sum(-points_xy * rel_xy, dim=-1)/(
            torch.norm(points_xy, 2, -1)*(torch.norm(rel_xy, 2, -1)))
        remove_index = ((point_cos >= 0.7) & (points_norm > self.key_distance))
        new_points = frame.points[~remove_index]
        new_cos = frame.get_pointsCos()[~remove_index]
        return LidarFrame(frame.index, new_points, new_cos,
                          frame.pose, new_keyframe=True)
//  [python]
def remove_back_points(self, frame):
        rel_pose = frame.get_rel_pose()
        points = frame.get_points()
        points_norm = torch.linalg.norm(points, 2, -1)
        points_xy = points[:, :2]
        if rel_pose == None:
            x = 1
            y = 0
        else:
            x = rel_pose[0, 3]
            y = rel_pose[1, 3]
        rel_xy = torch.ones((1, 2))
        rel_xy[0, 0] = x
        rel_xy[0, 1] = y
        point_cos = torch.sum(-points_xy * rel_xy, dim=-1)/(
            torch.linalg.norm(points_xy, 2, -1)*(torch.linalg.norm(rel_xy, 2, -1)))
        remove_index = ((point_cos >= 0.7) & (points_norm > self.key_distance))
        new_points = frame.points[~remove_index]
        new_cos = frame.get_pointsCos()[~remove_index]
        return LidarFrame(frame.index, new_points, new_cos,
                          frame.pose, new_keyframe=True)
//  [python]
def _compute_linear_displacement(
    coords: torch.Tensor, idxs: torch.Tensor
) -> torch.Tensor:
    """Computes the displacement of the BA and BC unit vectors in the linear angle
    "ABC". The displacements are measured along two axes that are perpendicular to the
    AC unit vector.

    Args:
        coords: The cartesian coordinates of a coords with shape=(n_atoms, 3)
        idxs: A tensor containing the indices of the atoms in each linear angle
            (first three columns) and the index of the axis to compute the displacement
            along (last column) with shape=(n_linear_angles, 4).
    Returns:
        A tensor of the linear displacements.
    """

    vector_ab = coords[idxs[:, 0]] - coords[idxs[:, 1]]
    vector_ab = vector_ab / torch.norm(vector_ab, dim=1).unsqueeze(1)

    vector_cb = coords[idxs[:, 2]] - coords[idxs[:, 1]]
    vector_cb = vector_cb / torch.norm(vector_cb, dim=1).unsqueeze(1)

    vector_ca = coords[idxs[:, 2]] - coords[idxs[:, 0]]
    vector_ca = vector_ca / torch.norm(vector_ca, dim=1).unsqueeze(1)

    # Take the dot product of each row of ``vector_ca`` with the x, y and z axis
    # and find the index (0 = x, 1 = y, 2 = z) of the axis that is most perpendicular
    # to each row. This ensures we don't try and take the cross-product of two
    # co-linear vectors.
    #
    # This is the same approach taken by geomeTRIC, but a bit more vectorized.
    basis = torch.tensor(
        [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], dtype=coords.dtype
    )
    basis_index = torch.argmin((vector_ca @ basis).square(), dim=-1)

    axis_0 = basis[basis_index]

    axis_1 = torch.cross(vector_ca, axis_0)
    axis_1 = axis_1 / torch.norm(axis_1, dim=1).unsqueeze(1)

    axis_2 = torch.cross(vector_ca, axis_1)
    axis_2 = axis_2 / torch.norm(axis_2, dim=1).unsqueeze(1)

    return torch.where(
        idxs[:, 3] == 0,
        (vector_ab * axis_1).sum(dim=-1) + (vector_cb * axis_1).sum(dim=-1),
        (vector_ab * axis_2).sum(dim=-1) + (vector_cb * axis_2).sum(dim=-1),
    )
//  [python]
def _compute_linear_displacement(
    coords: torch.Tensor, idxs: torch.Tensor
) -> torch.Tensor:
    """Computes the displacement of the BA and BC unit vectors in the linear angle
    "ABC". The displacements are measured along two axes that are perpendicular to the
    AC unit vector.

    Args:
        coords: The cartesian coordinates of a coords with shape=(n_atoms, 3)
        idxs: A tensor containing the indices of the atoms in each linear angle
            (first three columns) and the index of the axis to compute the displacement
            along (last column) with shape=(n_linear_angles, 4).
    Returns:
        A tensor of the linear displacements.
    """

    vector_ab = coords[idxs[:, 0]] - coords[idxs[:, 1]]
    vector_ab = vector_ab / torch.linalg.norm(vector_ab, dim=1, keepdim=True)

    vector_cb = coords[idxs[:, 2]] - coords[idxs[:, 1]]
    vector_cb = vector_cb / torch.linalg.norm(vector_cb, dim=1, keepdim=True)

    vector_ca = coords[idxs[:, 2]] - coords[idxs[:, 0]]
    vector_ca = vector_ca / torch.linalg.norm(vector_ca, dim=1, keepdim=True)

    # Take the dot product of each row of ``vector_ca`` with the x, y and z axis
    # and find the index (0 = x, 1 = y, 2 = z) of the axis that is most perpendicular
    # to each row. This ensures we don't try and take the cross-product of two
    # co-linear vectors.
    #
    # This is the same approach taken by geomeTRIC, but a bit more vectorized.
    basis = torch.tensor(
        [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], dtype=coords.dtype
    )
    basis_index = torch.argmin((vector_ca @ basis).square(), dim=-1)

    axis_0 = basis[basis_index]

    axis_1 = torch.cross(vector_ca, axis_0)
    axis_1 = axis_1 / torch.linalg.norm(axis_1, dim=1, keepdim=True)

    axis_2 = torch.cross(vector_ca, axis_1)
    axis_2 = axis_2 / torch.linalg.norm(axis_2, dim=1, keepdim=True)

    return torch.where(
        idxs[:, 3] == 0,
        (vector_ab * axis_1).sum(dim=-1) + (vector_cb * axis_1).sum(dim=-1),
        (vector_ab * axis_2).sum(dim=-1) + (vector_cb * axis_2).sum(dim=-1),
    )
//  [python]
def extract_energy(
        waveform: torch.Tensor, output_path: Optional[Path] = None,
        hop_length: int = 256, n_fft: int = 1024, log_scale: bool = True,
        phoneme_durations: Optional[List[int]] = None
):
    if output_path is not None and output_path.is_file():
        return

    assert len(waveform.shape) == 2 and waveform.shape[0] == 1
    waveform = waveform.view(1, 1, waveform.shape[1])
    waveform = F.pad(
        waveform.unsqueeze(1), [n_fft // 2, n_fft // 2, 0, 0],
        mode="reflect"
    )
    waveform = waveform.squeeze(1)

    fourier_basis = np.fft.fft(np.eye(n_fft))
    cutoff = int((n_fft / 2 + 1))
    fourier_basis = np.vstack(
        [np.real(fourier_basis[:cutoff, :]),
         np.imag(fourier_basis[:cutoff, :])]
    )

    forward_basis = torch.FloatTensor(fourier_basis[:, None, :])
    forward_transform = F.conv1d(
        waveform, forward_basis, stride=hop_length, padding=0
    )

    real_part = forward_transform[:, :cutoff, :]
    imag_part = forward_transform[:, cutoff:, :]
    magnitude = torch.sqrt(real_part ** 2 + imag_part ** 2)
    energy = torch.norm(magnitude, dim=1).squeeze(0).numpy()

    if phoneme_durations is not None:
        energy = trim_or_pad_to_target_length(energy, sum(phoneme_durations))
        d_cumsum = np.cumsum(np.concatenate([np.array([0]), phoneme_durations]))
        energy = np.array(
            [
                np.mean(energy[d_cumsum[i - 1]: d_cumsum[i]])
                for i in range(1, len(d_cumsum))
            ]
        )
        assert len(energy) == len(phoneme_durations)

    if log_scale:
        energy = np.log(energy + 1)

    if output_path is not None:
        np.save(output_path.as_posix(), energy)
    else:
        return energy
//  [python]
def extract_energy(
        waveform: torch.Tensor, output_path: Optional[Path] = None,
        hop_length: int = 256, n_fft: int = 1024, log_scale: bool = True,
        phoneme_durations: Optional[List[int]] = None
):
    if output_path is not None and output_path.is_file():
        return

    assert len(waveform.shape) == 2 and waveform.shape[0] == 1
    waveform = waveform.view(1, 1, waveform.shape[1])
    waveform = F.pad(
        waveform.unsqueeze(1), [n_fft // 2, n_fft // 2, 0, 0],
        mode="reflect"
    )
    waveform = waveform.squeeze(1)

    fourier_basis = np.fft.fft(np.eye(n_fft))
    cutoff = int((n_fft / 2 + 1))
    fourier_basis = np.vstack(
        [np.real(fourier_basis[:cutoff, :]),
         np.imag(fourier_basis[:cutoff, :])]
    )

    forward_basis = torch.FloatTensor(fourier_basis[:, None, :])
    forward_transform = F.conv1d(
        waveform, forward_basis, stride=hop_length, padding=0
    )

    real_part = forward_transform[:, :cutoff, :]
    imag_part = forward_transform[:, cutoff:, :]
    magnitude = torch.sqrt(real_part ** 2 + imag_part ** 2)
    energy = torch.linalg.norm(magnitude, dim=1).squeeze(0).numpy()

    if phoneme_durations is not None:
        energy = trim_or_pad_to_target_length(energy, sum(phoneme_durations))
        d_cumsum = np.cumsum(np.concatenate([np.array([0]), phoneme_durations]))
        energy = np.array(
            [
                np.mean(energy[d_cumsum[i - 1]: d_cumsum[i]])
                for i in range(1, len(d_cumsum))
            ]
        )
        assert len(energy) == len(phoneme_durations)

    if log_scale:
        energy = np.log(energy + 1)

    if output_path is not None:
        np.save(output_path.as_posix(), energy)
    else:
        return energy
//  [python]
def spectrogram(self, x):
        n_fft, hop_length, win_length = self.resolution
        x = F.pad(x, (int((n_fft - hop_length) / 2), int((n_fft - hop_length) / 2)), mode='reflect')
        x = x.squeeze(1).to(torch.float32)
        x = torch.stft(x, n_fft=n_fft, hop_length=hop_length, win_length=win_length, center=False) #[B, F, TT, 2]
        mag = torch.norm(x, p=2, dim =-1) #[B, F, TT]

        return mag
//  [python]
def spectrogram(self, x):
        n_fft, hop_length, win_length = self.resolution
        x = F.pad(x, (int((n_fft - hop_length) / 2), int((n_fft - hop_length) / 2)), mode='reflect')
        x = x.squeeze(1).to(torch.float32)
        x = torch.stft(x, n_fft=n_fft, hop_length=hop_length, win_length=win_length, center=False) #[B, F, TT, 2]
        mag = torch.linalg.norm(x, p=2, dim =-1) #[B, F, TT]

        return mag
//  [python]
def getRelTransError(c,T):
    diff = c - T
    err = torch.norm(diff,dim=1) / torch.norm(c,dim=1)
    return err.mean()
//  [python]
def getRelTransError(c,T):
    diff = c - T
    err = torch.linalg.norm(diff,dim=1) / torch.linalg.norm(c,dim=1)
    return err.mean()
//  [python]
def SC2_PCR(self, src_keypts, tgt_keypts):
        """
        Input:
            - src_keypts: [bs, num_corr, 3]
            - tgt_keypts: [bs, num_corr, 3]
        Output:
            - pred_trans:   [bs, 4, 4], the predicted transformation matrix.
            - pred_labels:  [bs, num_corr], the predicted inlier/outlier label (0,1), for classification loss calculation.
        """
        bs, num_corr = src_keypts.shape[0], tgt_keypts.shape[1]

        #################################
        # downsample points
        #################################
        if num_corr > self.max_points:
            src_keypts = src_keypts[:, :self.max_points, :]
            tgt_keypts = tgt_keypts[:, :self.max_points, :]
            num_corr = self.max_points

        #################################
        # compute cross dist
        #################################
        src_dist = torch.norm((src_keypts[:, :, None, :] - src_keypts[:, None, :, :]), dim=-1)
        target_dist = torch.norm((tgt_keypts[:, :, None, :] - tgt_keypts[:, None, :, :]), dim=-1)
        cross_dist = torch.abs(src_dist - target_dist)

        #################################
        # compute first order measure
        #################################
        SC_dist_thre = self.d_thre
        SC_measure = torch.clamp(1.0 - cross_dist ** 2 / SC_dist_thre ** 2, min=0)
        hard_SC_measure = (cross_dist < SC_dist_thre).float()

        #################################
        # select reliable seed correspondences
        #################################
        confidence = self.cal_leading_eigenvector(SC_measure, method='power')
        seeds = self.pick_seeds(src_dist, confidence, R=self.nms_radius, max_num=int(num_corr * self.ratio))

        #################################
        # compute second order measure
        #################################
        SC2_dist_thre = self.d_thre / 2
        hard_SC_measure_tight = (cross_dist < SC2_dist_thre).float()
        seed_hard_SC_measure = hard_SC_measure.gather(dim=1,
                                index=seeds[:, :, None].expand(-1, -1, num_corr))
        seed_hard_SC_measure_tight = hard_SC_measure_tight.gather(dim=1,
                                index=seeds[:, :, None].expand(-1, -1, num_corr))
        SC2_measure = torch.matmul(seed_hard_SC_measure_tight, hard_SC_measure_tight) * seed_hard_SC_measure

        #################################
        # compute the seed-wise transformations and select the best one
        #################################
        final_trans = self.cal_seed_trans(seeds, SC2_measure, src_keypts, tgt_keypts)

        #################################
        # refine the result by recomputing the transformation over the whole set
        #################################
        final_trans = self.post_refinement(final_trans, src_keypts, tgt_keypts, 20)

        return final_trans
//  [python]
def SC2_PCR(self, src_keypts, tgt_keypts):
        """
        Input:
            - src_keypts: [bs, num_corr, 3]
            - tgt_keypts: [bs, num_corr, 3]
        Output:
            - pred_trans:   [bs, 4, 4], the predicted transformation matrix.
            - pred_labels:  [bs, num_corr], the predicted inlier/outlier label (0,1), for classification loss calculation.
        """
        bs, num_corr = src_keypts.shape[0], tgt_keypts.shape[1]

        #################################
        # downsample points
        #################################
        if num_corr > self.max_points:
            src_keypts = src_keypts[:, :self.max_points, :]
            tgt_keypts = tgt_keypts[:, :self.max_points, :]
            num_corr = self.max_points

        #################################
        # compute cross dist
        #################################
        src_dist = torch.linalg.norm((src_keypts[:, :, None, :] - src_keypts[:, None, :, :]), dim=-1)
        target_dist = torch.linalg.norm((tgt_keypts[:, :, None, :] - tgt_keypts[:, None, :, :]), dim=-1)
        cross_dist = torch.abs(src_dist - target_dist)

        #################################
        # compute first order measure
        #################################
        SC_dist_thre = self.d_thre
        SC_measure = torch.clamp(1.0 - cross_dist ** 2 / SC_dist_thre ** 2, min=0)
        hard_SC_measure = (cross_dist < SC_dist_thre).float()

        #################################
        # select reliable seed correspondences
        #################################
        confidence = self.cal_leading_eigenvector(SC_measure, method='power')
        seeds = self.pick_seeds(src_dist, confidence, R=self.nms_radius, max_num=int(num_corr * self.ratio))

        #################################
        # compute second order measure
        #################################
        SC2_dist_thre = self.d_thre / 2
        hard_SC_measure_tight = (cross_dist < SC2_dist_thre).float()
        seed_hard_SC_measure = hard_SC_measure.gather(dim=1,
                                index=seeds[:, :, None].expand(-1, -1, num_corr))
        seed_hard_SC_measure_tight = hard_SC_measure_tight.gather(dim=1,
                                index=seeds[:, :, None].expand(-1, -1, num_corr))
        SC2_measure = torch.matmul(seed_hard_SC_measure_tight, hard_SC_measure_tight) * seed_hard_SC_measure

        #################################
        # compute the seed-wise transformations and select the best one
        #################################
        final_trans = self.cal_seed_trans(seeds, SC2_measure, src_keypts, tgt_keypts)

        #################################
        # refine the result by recomputing the transformation over the whole set
        #################################
        final_trans = self.post_refinement(final_trans, src_keypts, tgt_keypts, 20)

        return final_trans
//  [python]
def batch_rodrigues_torch(poses):
    """ poses: N x 3
    """
    batch_size = poses.shape[0]
    angle = torch.norm(poses + 1e-8, p=2, dim=1, keepdim=True)
    rot_dir = poses / angle

    cos = torch.cos(angle)[:, None]
    sin = torch.sin(angle)[:, None]

    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
    zeros = torch.zeros((batch_size, 1), device=poses.device)
    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1)
    K = K.reshape([batch_size, 3, 3])

    ident = torch.eye(3)[None].to(poses.device)
    rot_mat = ident + sin * K + (1 - cos) * torch.matmul(K, K)

    return rot_mat
//  [python]
def batch_rodrigues_torch(poses):
    """ poses: N x 3
    """
    batch_size = poses.shape[0]
    angle = torch.linalg.norm(poses + 1e-8, ord=2, dim=1, keepdim=True)
    rot_dir = poses / angle

    cos = torch.cos(angle)[:, None]
    sin = torch.sin(angle)[:, None]

    rx, ry, rz = torch.split(rot_dir, 1, dim=1)
    zeros = torch.zeros((batch_size, 1), device=poses.device)
    K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1)
    K = K.reshape([batch_size, 3, 3])

    ident = torch.eye(3)[None].to(poses.device)
    rot_mat = ident + sin * K + (1 - cos) * torch.matmul(K, K)

    return rot_mat
//  [python]
def getReprojErrorDLT(x_i,shape):
    m = x_i.shape[0]
    P = batchDLT(x_i, shape)
    s = torch.stack(m * [torch.cat((shape, torch.ones(shape.shape[0], 1)), dim=1).T])
    proj = torch.bmm(P, s)
    proj = proj / proj[:, 2].unsqueeze(1)
    diff = proj[:, :2, :] - x_i

    return torch.norm(diff,dim=1).mean()
//  [python]
def getReprojErrorDLT(x_i,shape):
    m = x_i.shape[0]
    P = batchDLT(x_i, shape)
    s = torch.stack(m * [torch.cat((shape, torch.ones(shape.shape[0], 1)), dim=1).T])
    proj = torch.bmm(P, s)
    proj = proj / proj[:, 2].unsqueeze(1)
    diff = proj[:, :2, :] - x_i

    return torch.linalg.norm(diff,dim=1).mean()
//  [python]
def compute_weight(self, module, do_power_iteration):
        weight = getattr(module, self.name + "_orig")
        u = getattr(module, self.name + "_u")
        v = getattr(module, self.name + "_v")
        curr_shape = weight.shape
        weight_mat = self.reshape_weight_to_matrix(weight)

        if do_power_iteration:
            with torch.no_grad():
                for _ in range(self.n_power_iterations):
                    # Spectral norm of weight equals to `u^T W v`,
                    # where `u` and `v`
                    # are the first left and right singular vectors.
                    # This power iteration produces approximations
                    # of `u` and `v`.
                    v = normalize(torch.mv(weight_mat.t(), u), dim=0, eps=self.eps, out=v)
                    u = normalize(torch.mv(weight_mat, v), dim=0, eps=self.eps, out=u)
                if self.n_power_iterations > 0:
                    # See above on why we need to clone
                    u = u.clone()
                    v = v.clone()

        sigma = torch.dot(u, torch.mv(weight_mat, v))
        if self.rank_ratio > 0.9999:
            return weight / sigma
        weight_mat = weight_mat / sigma

        rank_1 = torch.ger(u, v)
        residual = weight_mat - rank_1
        frob = torch.norm(residual)
        num = math.sqrt(self.rank - 1)
        if frob > num:
            weight_mat = residual * num / frob + rank_1
        weight_ = self.reshape_matrix_to_weight(weight_mat, curr_shape)
        return weight_
//  [python]
def compute_weight(self, module, do_power_iteration):
        weight = getattr(module, self.name + "_orig")
        u = getattr(module, self.name + "_u")
        v = getattr(module, self.name + "_v")
        curr_shape = weight.shape
        weight_mat = self.reshape_weight_to_matrix(weight)

        if do_power_iteration:
            with torch.no_grad():
                for _ in range(self.n_power_iterations):
                    # Spectral norm of weight equals to `u^T W v`,
                    # where `u` and `v`
                    # are the first left and right singular vectors.
                    # This power iteration produces approximations
                    # of `u` and `v`.
                    v = normalize(torch.mv(weight_mat.t(), u), dim=0, eps=self.eps, out=v)
                    u = normalize(torch.mv(weight_mat, v), dim=0, eps=self.eps, out=u)
                if self.n_power_iterations > 0:
                    # See above on why we need to clone
                    u = u.clone()
                    v = v.clone()

        sigma = torch.dot(u, torch.mv(weight_mat, v))
        if self.rank_ratio > 0.9999:
            return weight / sigma
        weight_mat = weight_mat / sigma

        rank_1 = torch.ger(u, v)
        residual = weight_mat - rank_1
        frob = torch.linalg.norm(residual)
        num = math.sqrt(self.rank - 1)
        if frob > num:
            weight_mat = residual * num / frob + rank_1
        weight_ = self.reshape_matrix_to_weight(weight_mat, curr_shape)
        return weight_
//  [python]
def _fast_h(v, stride=2):
    """
    Fast product of a series of Householder matrices. This implementation is oriented to the one introducesd in:
    https://invertibleworkshop.github.io/accepted_papers/pdfs/10.pdf
    This makes use of method 2 in: https://ecommons.cornell.edu/bitstream/handle/1813/6521/85-681.pdf?sequence=1&isAllowed=y

    :param v: Batched series of Householder matrices. The last dim is the dim of one vector and the second last is the
    number of elements in one product. This is the min amount of dims that need to be present.
    All further ones are considered batch dimensions.
    :param stride: Controls the number of parallel operations by the WY representation (see paper)
    should not be larger than half the number of matrices in one product.
    :return: The batched product of Householder matrices defined by v
    """
    assert v.ndim > 1
    assert stride <= v.shape[-2]

    d, m = v.shape[-2], v.shape[-1]
    k = d // stride
    last = k * stride
    v = v / torch.norm(v, dim=-1, p=2, keepdim=True)
    v = v.unsqueeze(-1)
    u = 2 * v
    ID = torch.eye(m, device=u.device)
    for dim in range(v.ndim-3):
        ID = ID.unsqueeze(0)

    # step 1 (compute intermediate groupings P_i)
    W = u[..., 0:last:stride, :, :]
    Y = v[..., 0:last:stride, :, :]

    for idx in range(1, stride):
        Pt = ID - torch.matmul(u[..., idx:last:stride, :, :], v[..., idx:last:stride, :, :].transpose(-1, -2))
        W = torch.cat([W, u[..., idx:last:stride, :, :]], dim=-1)
        Y = torch.cat([torch.matmul(Pt, Y), v[..., idx:last:stride, :, :]], dim=-1)

    # step 2 (multiply the WY reps)
    P = ID - torch.matmul(W[..., k-1, :, :], Y[..., k-1, :, :].transpose(-1, -2))
    for idx in reversed(range(0, k-1)):
        P = P - torch.matmul(W[..., idx, :, :], torch.matmul(Y[..., idx, :, :].transpose(-1, -2), P))

    # deal with the residual, using a stride of 2 here maxes the amount of parallel ops
    if d > last:
        even_end = d if (d-last) % 2 == 0 else d - 1
        W_resi = u[..., last:even_end:2, :, :]
        Y_resi = v[..., last:even_end:2, :, :]
        for idx in range(last+1, d if d == last+1 else last+2):
            Pt = ID - torch.matmul(u[..., idx:even_end:2, :, :], v[..., idx:even_end:2, :, :].transpose(-1, -2))
            W_resi = torch.cat([W_resi, u[..., idx:even_end:2, :, :]], dim=-1)
            Y_resi = torch.cat([torch.matmul(Pt, Y_resi), v[..., idx:even_end:2, :, :]], dim=-1)

        for idx in range(0, W_resi.shape[-3]):
            P = P - torch.matmul(P, torch.matmul(W_resi[..., idx, :, :], Y_resi[..., idx, :, :].transpose(-1, -2)))

        if even_end != d:
            P = P - torch.matmul(P, torch.matmul(u[..., -1, :, :], v[..., -1, :, :].transpose(-1, -2)))

    return P
//  [python]
def _fast_h(v, stride=2):
    """
    Fast product of a series of Householder matrices. This implementation is oriented to the one introducesd in:
    https://invertibleworkshop.github.io/accepted_papers/pdfs/10.pdf
    This makes use of method 2 in: https://ecommons.cornell.edu/bitstream/handle/1813/6521/85-681.pdf?sequence=1&isAllowed=y

    :param v: Batched series of Householder matrices. The last dim is the dim of one vector and the second last is the
    number of elements in one product. This is the min amount of dims that need to be present.
    All further ones are considered batch dimensions.
    :param stride: Controls the number of parallel operations by the WY representation (see paper)
    should not be larger than half the number of matrices in one product.
    :return: The batched product of Householder matrices defined by v
    """
    assert v.ndim > 1
    assert stride <= v.shape[-2]

    d, m = v.shape[-2], v.shape[-1]
    k = d // stride
    last = k * stride
    v = v / torch.linalg.norm(v, dim=-1, keepdim=True)
    v = v.unsqueeze(-1)
    u = 2 * v
    ID = torch.eye(m, device=u.device)
    for dim in range(v.ndim-3):
        ID = ID.unsqueeze(0)

    # step 1 (compute intermediate groupings P_i)
    W = u[..., 0:last:stride, :, :]
    Y = v[..., 0:last:stride, :, :]

    for idx in range(1, stride):
        Pt = ID - torch.matmul(u[..., idx:last:stride, :, :], v[..., idx:last:stride, :, :].transpose(-1, -2))
        W = torch.cat([W, u[..., idx:last:stride, :, :]], dim=-1)
        Y = torch.cat([torch.matmul(Pt, Y), v[..., idx:last:stride, :, :]], dim=-1)

    # step 2 (multiply the WY reps)
    P = ID - torch.matmul(W[..., k-1, :, :], Y[..., k-1, :, :].transpose(-1, -2))
    for idx in reversed(range(0, k-1)):
        P = P - torch.matmul(W[..., idx, :, :], torch.matmul(Y[..., idx, :, :].transpose(-1, -2), P))

    # deal with the residual, using a stride of 2 here maxes the amount of parallel ops
    if d > last:
        even_end = d if (d-last) % 2 == 0 else d - 1
        W_resi = u[..., last:even_end:2, :, :]
        Y_resi = v[..., last:even_end:2, :, :]
        for idx in range(last+1, d if d == last+1 else last+2):
            Pt = ID - torch.matmul(u[..., idx:even_end:2, :, :], v[..., idx:even_end:2, :, :].transpose(-1, -2))
            W_resi = torch.cat([W_resi, u[..., idx:even_end:2, :, :]], dim=-1)
            Y_resi = torch.cat([torch.matmul(Pt, Y_resi), v[..., idx:even_end:2, :, :]], dim=-1)

        for idx in range(0, W_resi.shape[-3]):
            P = P - torch.matmul(P, torch.matmul(W_resi[..., idx, :, :], Y_resi[..., idx, :, :].transpose(-1, -2)))

        if even_end != d:
            P = P - torch.matmul(P, torch.matmul(u[..., -1, :, :], v[..., -1, :, :].transpose(-1, -2)))

    return P
//  [python]
def getDistances(control_w):
    p1 = control_w[:3,0]
    p2 = control_w[:3,1]
    p3 = control_w[:3,2]
    p4 = control_w[:3,3]

    d12 = torch.norm(p1-p2,p=2)
    d13 = torch.norm(p1-p3,p=2)
    d14 = torch.norm(p1-p4,p=2)
    d23 = torch.norm(p2-p3,p=2)
    d24 = torch.norm(p2-p4,p=2)
    d34 = torch.norm(p3-p4,p=2)

    return d12,d13,d14,d23,d24,d34