jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def broyden(g, x_init, J_inv_init, max_steps=50, cvg_thresh=1e-5, dvg_thresh=1, eps=1e-6):
    """Find roots of the given function g(x) = 0.
    This function is impleneted based on https://github.com/locuslab/deq.

    Tensor shape abbreviation:
        N: number of points
        D: space dimension
    Args:
        g (function): the function of which the roots are to be determined. shape: [N, D, 1]->[N, D, 1]
        x_init (tensor): initial value of the parameters. shape: [N, D, 1]
        J_inv_init (tensor): initial value of the inverse Jacobians. shape: [N, D, D]

        max_steps (int, optional): max number of iterations. Defaults to 50.
        cvg_thresh (float, optional): covergence threshold. Defaults to 1e-5.
        dvg_thresh (float, optional): divergence threshold. Defaults to 1.
        eps (float, optional): a small number added to the denominator to prevent numerical error. Defaults to 1e-6.

    Returns:
        result (tensor): root of the given function. shape: [N, D, 1]
        diff (tensor): corresponding loss. [N]
        valid_ids (tensor): identifiers of converged points. [N]
    """

    # initialization
    x = x_init.clone().detach()
    J_inv = J_inv_init.clone().detach()

    ids_val = torch.ones(x.shape[0]).bool()

    gx = g(x, mask=ids_val)
    update = -J_inv.bmm(gx)

    x_opt = x
    gx_norm_opt = torch.norm(gx.squeeze(-1), dim=-1)

    delta_gx = torch.zeros_like(gx)
    delta_x = torch.zeros_like(x)

    ids_val = torch.ones_like(gx_norm_opt).bool()

    for _ in range(max_steps):

        # update paramter values
        delta_x[ids_val] = update
        x[ids_val] += delta_x[ids_val]
        delta_gx[ids_val] = g(x, mask=ids_val) - gx[ids_val]
        gx[ids_val] += delta_gx[ids_val]

        # store values with minial loss
        gx_norm = torch.norm(gx.squeeze(-1), dim=-1)
        ids_opt = gx_norm < gx_norm_opt
        gx_norm_opt[ids_opt] = gx_norm.clone().detach()[ids_opt]
        x_opt[ids_opt] = x.clone().detach()[ids_opt]

        # exclude converged and diverged points from furture iterations
        ids_val = (gx_norm_opt > cvg_thresh) & (gx_norm < dvg_thresh)
        if ids_val.sum() <= 0:
            break

        # compute paramter update for next iter
        vT = (delta_x[ids_val]).transpose(-1, -2).bmm(J_inv[ids_val])
        a = delta_x[ids_val] - J_inv[ids_val].bmm(delta_gx[ids_val])
        b = vT.bmm(delta_gx[ids_val])
        b[b >= 0] += eps
        b[b < 0] -= eps
        u = a / b
        J_inv[ids_val] += u.bmm(vT)
        update = -J_inv[ids_val].bmm(gx[ids_val])

    return {'result': x_opt, 'diff': gx_norm_opt, 'valid_ids': gx_norm_opt < cvg_thresh}
//  [python]
def broyden(g, x_init, J_inv_init, max_steps=50, cvg_thresh=1e-5, dvg_thresh=1, eps=1e-6):
    """Find roots of the given function g(x) = 0.
    This function is impleneted based on https://github.com/locuslab/deq.

    Tensor shape abbreviation:
        N: number of points
        D: space dimension
    Args:
        g (function): the function of which the roots are to be determined. shape: [N, D, 1]->[N, D, 1]
        x_init (tensor): initial value of the parameters. shape: [N, D, 1]
        J_inv_init (tensor): initial value of the inverse Jacobians. shape: [N, D, D]

        max_steps (int, optional): max number of iterations. Defaults to 50.
        cvg_thresh (float, optional): covergence threshold. Defaults to 1e-5.
        dvg_thresh (float, optional): divergence threshold. Defaults to 1.
        eps (float, optional): a small number added to the denominator to prevent numerical error. Defaults to 1e-6.

    Returns:
        result (tensor): root of the given function. shape: [N, D, 1]
        diff (tensor): corresponding loss. [N]
        valid_ids (tensor): identifiers of converged points. [N]
    """

    # initialization
    x = x_init.clone().detach()
    J_inv = J_inv_init.clone().detach()

    ids_val = torch.ones(x.shape[0]).bool()

    gx = g(x, mask=ids_val)
    update = -J_inv.bmm(gx)

    x_opt = x
    gx_norm_opt = torch.linalg.norm(gx.squeeze(-1), dim=-1)

    delta_gx = torch.zeros_like(gx)
    delta_x = torch.zeros_like(x)

    ids_val = torch.ones_like(gx_norm_opt).bool()

    for _ in range(max_steps):

        # update paramter values
        delta_x[ids_val] = update
        x[ids_val] += delta_x[ids_val]
        delta_gx[ids_val] = g(x, mask=ids_val) - gx[ids_val]
        gx[ids_val] += delta_gx[ids_val]

        # store values with minial loss
        gx_norm = torch.linalg.norm(gx.squeeze(-1), dim=-1)
        ids_opt = gx_norm < gx_norm_opt
        gx_norm_opt[ids_opt] = gx_norm.clone().detach()[ids_opt]
        x_opt[ids_opt] = x.clone().detach()[ids_opt]

        # exclude converged and diverged points from furture iterations
        ids_val = (gx_norm_opt > cvg_thresh) & (gx_norm < dvg_thresh)
        if ids_val.sum() <= 0:
            break

        # compute paramter update for next iter
        vT = (delta_x[ids_val]).transpose(-1, -2).bmm(J_inv[ids_val])
        a = delta_x[ids_val] - J_inv[ids_val].bmm(delta_gx[ids_val])
        b = vT.bmm(delta_gx[ids_val])
        b[b >= 0] += eps
        b[b < 0] -= eps
        u = a / b
        J_inv[ids_val] += u.bmm(vT)
        update = -J_inv[ids_val].bmm(gx[ids_val])

    return {'result': x_opt, 'diff': gx_norm_opt, 'valid_ids': gx_norm_opt < cvg_thresh}
//  [python]
def rho_criterion(self, c: torch.Tensor = None):
        differences = torch.stack(
            [
                c[:, j : j + 1] - c[:, j : j + 1].T * torch.ones(self.k, self.k)
                for j in range(self.latent_dim)
            ],
            dim=2,
        )

        differences_norm = torch.norm(differences, dim=-1) + 1e16 * torch.eye(
            self.k
        )

        min_dist_between_neigh = torch.min(differences_norm, dim=-1).values

        rho = torch.max(min_dist_between_neigh)

        return rho
//  [python]
def rho_criterion(self, c: torch.Tensor = None):
        differences = torch.stack(
            [
                c[:, j : j + 1] - c[:, j : j + 1].T * torch.ones(self.k, self.k)
                for j in range(self.latent_dim)
            ],
            dim=2,
        )

        differences_norm = torch.linalg.norm(differences, dim=-1) + 1e16 * torch.eye(
            self.k
        )

        min_dist_between_neigh = torch.min(differences_norm, dim=-1).values

        rho = torch.max(min_dist_between_neigh)

        return rho
//  [python]
def G(self, z: torch.Tensor = None) -> torch.Tensor:
        if z.requires_grad == False:
            z.requires_grad = True

        g_z = sum(
            [cov * basis(z=z) for cov, basis in zip(self.covariances, self.basis)]
        ) + self.lambd * torch.exp(-self.tau * torch.norm(z, 2)) * torch.eye(
            self.latent_dim
        )

        return g_z
//  [python]
def G(self, z: torch.Tensor = None) -> torch.Tensor:
        if z.requires_grad == False:
            z.requires_grad = True

        g_z = sum(
            [cov * basis(z=z) for cov, basis in zip(self.covariances, self.basis)]
        ) + self.lambd * torch.exp(-self.tau * torch.linalg.norm(z, 2)) * torch.eye(
            self.latent_dim
        )

        return g_z
//  [python]
def test_linalg_norm() -> None:
    x: Tensor[torch.float32, L[2], L[3], L[4]]
    y: Tensor[torch.float32, L[2]] = torch.norm(x, dim=(-2, -1))
    # pyre-fixme[9]: Expected error.
    y_error: Tensor[torch.float32, L[99]] = torch.norm(x, dim=(-2, -1))
//  [python]
def test_linalg_norm() -> None:
    x: Tensor[torch.float32, L[2], L[3], L[4]]
    y: Tensor[torch.float32, L[2]] = torch.linalg.norm(x, dim=(-2, -1))
    # pyre-fixme[9]: Expected error.
    y_error: Tensor[torch.float32, L[99]] = torch.linalg.norm(x, dim=(-2, -1))
//  [python]
def dist_constrain(self, inbox_inds, cluster_xyz, gt_bboxes_3d, gt_labels_3d):

        inbox_inds = inbox_inds.clone()
        max_dist = self.train_cfg.get('max_assign_dist', None)
        if max_dist is None:
            return inbox_inds

        if not (inbox_inds > -1).any():
            return inbox_inds

        pos_mask = inbox_inds > -1
        pos_inds = inbox_inds[pos_mask].clone()
        pos_xyz = cluster_xyz[pos_mask]
        pos_labels = gt_labels_3d[pos_inds]
        pos_box_center = gt_bboxes_3d.gravity_center[pos_inds]
        rel_dist = torch.norm(pos_xyz[:, :2] - pos_box_center[:, :2], p=2, dim=1) # only xy-dist
        thresh = torch.zeros_like(rel_dist)
        assert len(max_dist) == self.num_classes
        for i in range(self.num_classes):
            thresh[pos_labels == i] = max_dist[i]
        
        pos_inds[rel_dist > thresh] = -1
        inbox_inds[pos_mask] = pos_inds
        return inbox_inds
//  [python]
def dist_constrain(self, inbox_inds, cluster_xyz, gt_bboxes_3d, gt_labels_3d):

        inbox_inds = inbox_inds.clone()
        max_dist = self.train_cfg.get('max_assign_dist', None)
        if max_dist is None:
            return inbox_inds

        if not (inbox_inds > -1).any():
            return inbox_inds

        pos_mask = inbox_inds > -1
        pos_inds = inbox_inds[pos_mask].clone()
        pos_xyz = cluster_xyz[pos_mask]
        pos_labels = gt_labels_3d[pos_inds]
        pos_box_center = gt_bboxes_3d.gravity_center[pos_inds]
        rel_dist = torch.linalg.norm(pos_xyz[:, :2] - pos_box_center[:, :2], ord=2, dim=1) # only xy-dist
        thresh = torch.zeros_like(rel_dist)
        assert len(max_dist) == self.num_classes
        for i in range(self.num_classes):
            thresh[pos_labels == i] = max_dist[i]
        
        pos_inds[rel_dist > thresh] = -1
        inbox_inds[pos_mask] = pos_inds
        return inbox_inds
//  [python]
def is_feasible(self, x, atol=1e-5, rtol=1e-5):
        norms = torch.norm(x, dim=(-2, -1), p='nuc')
        return (norms <= self.alpha * (1. + rtol) + atol)
//  [python]
def is_feasible(self, x, atol=1e-5, rtol=1e-5):
        norms = torch.linalg.norm(x, dim=(-2, -1), ord='nuc')
        return (norms <= self.alpha * (1. + rtol) + atol)
//  [python]
def get_group_norms(self, x):
        """Compute the vector of L2 norms within groups"""
        group_norms = []
        for g in self.groups:
            subtensor = x[g]

            group_norms.append(torch.norm(subtensor, dim=-1))

        group_norms = torch.stack(group_norms, dim=-1)
        return group_norms
//  [python]
def get_group_norms(self, x):
        """Compute the vector of L2 norms within groups"""
        group_norms = []
        for g in self.groups:
            subtensor = x[g]

            group_norms.append(torch.linalg.norm(subtensor, dim=-1))

        group_norms = torch.stack(group_norms, dim=-1)
        return group_norms
//  [python]
def forward(self, x):
        x = self.linear(x)
        x = self.relu(x)
        x = torch.norm(x, p=2, dim=1)
        x = self.relu(x)
        return x
//  [python]
def forward(self, x):
        x = self.linear(x)
        x = self.relu(x)
        x = torch.linalg.norm(x, ord=2, dim=1)
        x = self.relu(x)
        return x
//  [python]
def is_feasible(self, x, rtol=1e-5, atol=1e-7):
        group_norms = self.get_group_norms(x)
        return torch.norm(group_norms, p=1, dim=-1) <= (self.alpha * (1. + rtol)
                                                                 + atol)
//  [python]
def is_feasible(self, x, rtol=1e-5, atol=1e-7):
        group_norms = self.get_group_norms(x)
        return torch.linalg.norm(group_norms, ord=1, dim=-1) <= (self.alpha * (1. + rtol)
                                                                 + atol)
//  [python]
def angle_between(self, sph_v1, sph_v2):
        def sph2cart(sv):
            r, theta, phi = sv[0], sv[1], sv[2]
            return torch.tensor([r * torch.sin(theta) * torch.cos(phi), r * torch.sin(theta) * torch.sin(phi), r * torch.cos(theta)])
        def unit_vector(v):
            return v / torch.norm(v)
        def angle_between_2_sph(sv1, sv2):
            v1, v2 = sph2cart(sv1), sph2cart(sv2)
            v1_u, v2_u = unit_vector(v1), unit_vector(v2)
            return torch.arccos(torch.clip(torch.dot(v1_u, v2_u), -1.0, 1.0))
        angles = torch.empty(len(sph_v1), len(sph_v2))
        for i, sv1 in enumerate(sph_v1):
            for j, sv2 in enumerate(sph_v2):
                angles[i][j] = angle_between_2_sph(sv1, sv2)
        return angles
//  [python]
def angle_between(self, sph_v1, sph_v2):
        def sph2cart(sv):
            r, theta, phi = sv[0], sv[1], sv[2]
            return torch.tensor([r * torch.sin(theta) * torch.cos(phi), r * torch.sin(theta) * torch.sin(phi), r * torch.cos(theta)])
        def unit_vector(v):
            return v / torch.linalg.norm(v)
        def angle_between_2_sph(sv1, sv2):
            v1, v2 = sph2cart(sv1), sph2cart(sv2)
            v1_u, v2_u = unit_vector(v1), unit_vector(v2)
            return torch.arccos(torch.clip(torch.dot(v1_u, v2_u), -1.0, 1.0))
        angles = torch.empty(len(sph_v1), len(sph_v2))
        for i, sv1 in enumerate(sph_v1):
            for j, sv2 in enumerate(sph_v2):
                angles[i][j] = angle_between_2_sph(sv1, sv2)
        return angles
//  [python]
def sample_sphere_points(n_points, radius=1.0):
    points = torch.randn(n_points, 3)
    points /= torch.norm(points, dim=1, keepdim=True)
    points *= radius
    return points
//  [python]
def sample_sphere_points(n_points, radius=1.0):
    points = torch.randn(n_points, 3)
    points /= torch.linalg.norm(points, axis=1).view(-1, 1)
    points *= radius
    return points
//  [python]
def batch_normalized_mse(pred: Tensor, target: Tensor) -> Tensor:
    """Calculates batch-wise normalized mse error between two tensors."""

    pred_flat = pred.reshape(pred.size(0), -1)
    target_flat = target.reshape(target.size(0), -1)

    diff_norms = torch.norm(pred_flat - target_flat, p=2, dim=1)
    target_norms = torch.norm(target_flat, p=2, dim=1)

    error = diff_norms / target_norms
    return torch.mean(error)
//  [python]
def batch_normalized_mse(pred: Tensor, target: Tensor) -> Tensor:
    """Calculates batch-wise normalized mse error between two tensors."""

    pred_flat = pred.reshape(pred.size(0), -1)
    target_flat = target.reshape(target.size(0), -1)

    diff_norms = torch.linalg.norm(pred_flat - target_flat, ord=2.0, dim=1)
    target_norms = torch.linalg.norm(target_flat, ord=2.0, dim=1)

    error = diff_norms / target_norms
    return torch.mean(error)
//  [python]
def test_model_normalization(self):
        names = self.large_clip_models + self.e5_models
        sentences = ['hello', 'this is a test sentence. so is this.', ['hello', 'this is a test sentence. so is this.']]
        device = "cuda"
        eps = 1e-6

        for name in names:
            model_properties = get_model_properties_from_registry(name)
            model = _load_model(model_properties['name'], model_properties=model_properties, device=device)

            for sentence in sentences:
                output = model.encode(sentence, normalize=True)
                output = _convert_vectorized_output(output)
                max_output_norm = max(torch.norm(FloatTensor(output), dim=1))
                min_output_norm = min(torch.norm(FloatTensor(output), dim=1))

                assert abs(max_output_norm - 1) < eps, f"{name}, {sentence}"
                assert abs(min_output_norm - 1) < eps, f"{name}, {sentence}"

            del model
            clear_loaded_models()
//  [python]
def test_model_normalization(self):
        names = self.large_clip_models + self.e5_models
        sentences = ['hello', 'this is a test sentence. so is this.', ['hello', 'this is a test sentence. so is this.']]
        device = "cuda"
        eps = 1e-6

        for name in names:
            model_properties = get_model_properties_from_registry(name)
            model = _load_model(model_properties['name'], model_properties=model_properties, device=device)

            for sentence in sentences:
                output = model.encode(sentence, normalize=True)
                output = _convert_vectorized_output(output)
                max_output_norm = max(torch.linalg.norm(FloatTensor(output), dim=1))
                min_output_norm = min(torch.linalg.norm(FloatTensor(output), dim=1))

                assert abs(max_output_norm - 1) < eps, f"{name}, {sentence}"
                assert abs(min_output_norm - 1) < eps, f"{name}, {sentence}"

            del model
            clear_loaded_models()
//  [python]
def _evaluate_batch(self, batch: SolutionBatch):
        values = torch.as_tensor(batch.access_values(), dtype=torch.float32)
        evals = batch.access_evals()
        evals[:, 0] = torch.norm(values, dim=-1)
        evals[:, 1] = torch.sin(torch.sum(values, dim=-1))
        evals[:, 2] = torch.cos(torch.sum(values, dim=-1))
//  [python]
def _evaluate_batch(self, batch: SolutionBatch):
        values = torch.as_tensor(batch.access_values(), dtype=torch.float32)
        evals = batch.access_evals()
        evals[:, 0] = torch.linalg.norm(values, dim=-1)
        evals[:, 1] = torch.sin(torch.sum(values, dim=-1))
        evals[:, 2] = torch.cos(torch.sum(values, dim=-1))
//  [python]
def contract_samples(x, order=float('inf')):
    mag = t.norm(x, order, dim=-1)[..., None]
    return torch.where(mag < 1, x, (2 - (1 / mag)) * (x / mag)), mag
//  [python]
def contract_samples(x, order=float('inf')):
    mag = LA.norm(x, order, dim=-1)[..., None]
    return torch.where(mag < 1, x, (2 - (1 / mag)) * (x / mag)), mag
//  [python]
def test_error(module, hidden_states, target_states, cache, attn_params):

    rfn_sum = 0
    rfn_count = 0
    for x, xref in zip(hidden_states, target_states):
        x = x.cuda()
        xref = xref.cuda()
        xtest = module.forward(x, cache, attn_params)
        xtest = xtest[0].float()
        xref = xref[0].float()
        rfn_sum += (torch.norm(xtest - xref) / torch.norm(xref)).item()
        rfn_count += 1

    return max(1e-6, 1 - (rfn_sum / rfn_count))
//  [python]
def test_error(module, hidden_states, target_states, cache, attn_params):

    rfn_sum = 0
    rfn_count = 0
    for x, xref in zip(hidden_states, target_states):
        x = x.cuda()
        xref = xref.cuda()
        xtest = module.forward(x, cache, attn_params)
        xtest = xtest[0].float()
        xref = xref[0].float()
        rfn_sum += (torch.linalg.norm(xtest - xref, 'fro') / torch.linalg.norm(xref, 'fro')).item()
        rfn_count += 1

    return max(1e-6, 1 - (rfn_sum / rfn_count))
//  [python]
def _test_matrix_to_quaternion():
    # init a random batch of quaternion
    r = torch.randn((10, 4)).cuda()

    norm = torch.sqrt(
        r[:, 0] * r[:, 0] + r[:, 1] * r[:, 1] + r[:, 2] * r[:, 2] + r[:, 3] * r[:, 3]
    )

    q = r / norm[:, None]

    q = standardize_quaternion(q)

    R = quternion_to_matrix(q)

    I_rec = R @ R.transpose(-2, -1)
    I_rec_error = torch.abs(I_rec - torch.eye(3, device="cuda")[None, ...]).max()

    q_recovered = matrix_to_quaternion(R)
    norm_ = torch.norm(q_recovered, dim=-1)
    q_recovered = q_recovered / norm_[..., None]
    q_recovered = standardize_quaternion(q_recovered)

    print(q_recovered.shape, q.shape, R.shape)

    rec = (q - q_recovered).abs().max()

    print("rotation to I error:", I_rec_error, "quant rec error: ", rec)
//  [python]
def _test_matrix_to_quaternion():
    # init a random batch of quaternion
    r = torch.randn((10, 4)).cuda()

    norm = torch.sqrt(
        r[:, 0] * r[:, 0] + r[:, 1] * r[:, 1] + r[:, 2] * r[:, 2] + r[:, 3] * r[:, 3]
    )

    q = r / norm[:, None]

    q = standardize_quaternion(q)

    R = quternion_to_matrix(q)

    I_rec = R @ R.transpose(-2, -1)
    I_rec_error = torch.abs(I_rec - torch.eye(3, device="cuda")[None, ...]).max()

    q_recovered = matrix_to_quaternion(R)
    norm_ = torch.linalg.norm(q_recovered, dim=-1)
    q_recovered = q_recovered / norm_[..., None]
    q_recovered = standardize_quaternion(q_recovered)

    print(q_recovered.shape, q.shape, R.shape)

    rec = (q - q_recovered).abs().max()

    print("rotation to I error:", I_rec_error, "quant rec error: ", rec)
//  [python]
def inverse_contract_samples(x, mag_origial,order=float('inf')):
    mag = t.norm(x, order, dim=-1)[..., None]
    return torch.where(mag < 1, x, (x*mag_origial)/(2-(1/mag_origial)))
//  [python]
def inverse_contract_samples(x, mag_origial,order=float('inf')):
    mag = LA.norm(x, order, dim=-1)[..., None]
    return torch.where(mag < 1, x, (x*mag_origial)/(2-(1/mag_origial)))
//  [python]
def project_M(M, M_constraint):
	if M_constraint == 'simplex':
		project2simplex(M, dim=0)
	elif M_constraint == 'unit sphere':
		M.div_(torch.norm(M, p=2, dim=0, keepdim=True))
	elif M_constraint == 'nonneg unit sphere':
		M.clip_(1e-10).div_(torch.norm(M, p=2, dim=0, keepdim=True))
	else:
		raise NotImplementedError()
	return M
//  [python]
def project_M(M, M_constraint):
	if M_constraint == 'simplex':
		project2simplex(M, dim=0)
	elif M_constraint == 'unit sphere':
		M.div_(torch.linalg.norm(M, ord=2, dim=0, keepdim=True))
	elif M_constraint == 'nonneg unit sphere':
		M.clip_(1e-10).div_(torch.linalg.norm(M, ord=2, dim=0, keepdim=True))
	else:
		raise NotImplementedError()
	return M
//  [python]
def pos_map_to_mesh(pos_map: torch.Tensor):
    pd = (0, 0, 1 if pos_map.shape[1] % 2 == 1 else 0, 0, 1 if pos_map.shape[0] % 2 == 1 else 0, 0)
    pos_map = F.pad(pos_map, pd, 'constant', 0)
    mask = torch.norm(pos_map, dim = -1) > 0.1
    # cv.imshow('mask', mask.cpu().numpy().astype(np.uint8) * 255)
    mask = cv.erode(mask.cpu().numpy().astype(np.uint8), (5, 5), iterations = 20)
    mask = torch.from_numpy(mask > 0).to(pos_map.device)
    # cv.imshow('mask_eroded', mask.cpu().numpy().astype(np.uint8) * 255)
    # cv.waitKey(0)
    v0 = pos_map[:-1, :-1].reshape(-1, 3)
    v1 = pos_map[1:, :-1].reshape(-1, 3)
    v2 = pos_map[:-1, 1:].reshape(-1, 3)
    v3 = pos_map[1:, 1:].reshape(-1, 3)
    m0 = mask[:-1, :-1].reshape(-1)
    m1 = mask[1:, :-1].reshape(-1)
    m2 = mask[:-1, 1:].reshape(-1)
    m3 = mask[1:, 1:].reshape(-1)
    vertices = torch.cat([v0, v1, v2, v3], 0)
    masks = torch.cat([m0, m1, m2, m3], 0)
    pnum = v0.shape[0]

    a = torch.arange(0, pnum).to(torch.int64).to(pos_map.device)
    f1 = torch.stack([a, a + pnum, a + pnum * 2], 1)
    f2 = torch.stack([a + pnum, a + pnum * 3, a + pnum * 2], 1)
    faces = torch.cat([f1, f2], 0)

    # remove invalid faces
    face_mask = masks[faces.reshape(-1)].reshape(-1, 3).sum(1) == 3
    face_mask = torch.logical_and(face_mask, torch.norm(vertices[faces[:, 0]] - vertices[faces[:, 1]], dim = 1) < 0.02)
    face_mask = torch.logical_and(face_mask, torch.norm(vertices[faces[:, 1]] - vertices[faces[:, 2]], dim = 1) < 0.02)
    face_mask = torch.logical_and(face_mask, torch.norm(vertices[faces[:, 0]] - vertices[faces[:, 2]], dim = 1) < 0.02)
    valid_faces = faces[face_mask]

    # debug
    mesh = trimesh.Trimesh(vertices = vertices.cpu().numpy(), faces = valid_faces.cpu().numpy())
    mesh.export('./debug/mesh.obj')
    exit(1)
//  [python]
def pos_map_to_mesh(pos_map: torch.Tensor):
    pd = (0, 0, 1 if pos_map.shape[1] % 2 == 1 else 0, 0, 1 if pos_map.shape[0] % 2 == 1 else 0, 0)
    pos_map = F.pad(pos_map, pd, 'constant', 0)
    mask = torch.linalg.norm(pos_map, dim = -1) > 0.1
    # cv.imshow('mask', mask.cpu().numpy().astype(np.uint8) * 255)
    mask = cv.erode(mask.cpu().numpy().astype(np.uint8), (5, 5), iterations = 20)
    mask = torch.from_numpy(mask > 0).to(pos_map.device)
    # cv.imshow('mask_eroded', mask.cpu().numpy().astype(np.uint8) * 255)
    # cv.waitKey(0)
    v0 = pos_map[:-1, :-1].reshape(-1, 3)
    v1 = pos_map[1:, :-1].reshape(-1, 3)
    v2 = pos_map[:-1, 1:].reshape(-1, 3)
    v3 = pos_map[1:, 1:].reshape(-1, 3)
    m0 = mask[:-1, :-1].reshape(-1)
    m1 = mask[1:, :-1].reshape(-1)
    m2 = mask[:-1, 1:].reshape(-1)
    m3 = mask[1:, 1:].reshape(-1)
    vertices = torch.cat([v0, v1, v2, v3], 0)
    masks = torch.cat([m0, m1, m2, m3], 0)
    pnum = v0.shape[0]

    a = torch.arange(0, pnum).to(torch.int64).to(pos_map.device)
    f1 = torch.stack([a, a + pnum, a + pnum * 2], 1)
    f2 = torch.stack([a + pnum, a + pnum * 3, a + pnum * 2], 1)
    faces = torch.cat([f1, f2], 0)

    # remove invalid faces
    face_mask = masks[faces.reshape(-1)].reshape(-1, 3).sum(1) == 3
    face_mask = torch.logical_and(face_mask, torch.linalg.norm(vertices[faces[:, 0]] - vertices[faces[:, 1]], dim = 1) < 0.02)
    face_mask = torch.logical_and(face_mask, torch.linalg.norm(vertices[faces[:, 1]] - vertices[faces[:, 2]], dim = 1) < 0.02)
    face_mask = torch.logical_and(face_mask, torch.linalg.norm(vertices[faces[:, 0]] - vertices[faces[:, 2]], dim = 1) < 0.02)
    valid_faces = faces[face_mask]

    # debug
    mesh = trimesh.Trimesh(vertices = vertices.cpu().numpy(), faces = valid_faces.cpu().numpy())
    mesh.export('./debug/mesh.obj')
    exit(1)
//  [python]
def test_joint_transform_parse(usd_cuda_robot, urdf_cuda_robot):
    usd_pose = usd_cuda_robot.get_all_link_transforms()
    urdf_pose = urdf_cuda_robot.get_all_link_transforms()
    p, q = usd_pose.distance(urdf_pose)
    p = torch.norm(q)
    q = torch.norm(q)
    assert p < 1e-5 and q < 1e-5
//  [python]
def test_joint_transform_parse(usd_cuda_robot, urdf_cuda_robot):
    usd_pose = usd_cuda_robot.get_all_link_transforms()
    urdf_pose = urdf_cuda_robot.get_all_link_transforms()
    p, q = usd_pose.distance(urdf_pose)
    p = torch.linalg.norm(q)
    q = torch.linalg.norm(q)
    assert p < 1e-5 and q < 1e-5
//  [python]
def predict_pos_score(self, latents, num_atoms, atom_types, lengths, angles, cart_coords, batch, sigma):
        edge_index, distance_vectors, _ = get_pbc_cutoff_graphs(cart_coords, lengths, angles, num_atoms, self.cutoff, self.max_num_neighbors)
        edge_features = self.backbone(atom_types, edge_index, distance_vectors)
        num_multi_edge_per_graph = scatter(torch.ones(size=(edge_index.shape[1],), device=edge_index.device).long(), batch[edge_index[0]], reduce='sum')
        latents_per_multi_edge = latents.repeat_interleave(num_multi_edge_per_graph, dim=0)
        edge_features = torch.cat((edge_features, latents_per_multi_edge), dim=1)

        j, i = edge_index
        no_iden_mask = (i != j)
        j, i, edge_features, distance_vectors = j[no_iden_mask], i[no_iden_mask], edge_features[no_iden_mask], distance_vectors[no_iden_mask]
        dists_per_multi_edge = torch.norm(distance_vectors, dim=-1, keepdim=True)
        scores_per_multi_edge = self.fc_score(edge_features) * distance_vectors / dists_per_multi_edge

        num_multi_edges = len(i)
        new_edge_start_mask = torch.logical_or(i[:-1] != i[1:], j[:-1] != j[1:])
        new_edge_start_id = torch.nonzero(new_edge_start_mask).view(-1) + 1
        num_multi_edges_per_edge = torch.cat([new_edge_start_id[0:1], new_edge_start_id[1:] - new_edge_start_id[:-1], num_multi_edges - new_edge_start_id[-1:]])
        multi_edge_to_edge_idx = torch.repeat_interleave(torch.arange(len(num_multi_edges_per_edge), device=num_multi_edges_per_edge.device), num_multi_edges_per_edge)
        scores_per_edge = scatter(scores_per_multi_edge, multi_edge_to_edge_idx, dim=0, reduce='mean')
        unique_edge_receiver_index = scatter(i, multi_edge_to_edge_idx, dim=0, reduce='mean').long()
        scores_per_node_pos = scatter(scores_per_edge, unique_edge_receiver_index, dim=0, dim_size=len(batch), reduce='sum')

        score_norm = self.get_score_norm(sigma)
        return scores_per_node_pos / score_norm
//  [python]
def predict_pos_score(self, latents, num_atoms, atom_types, lengths, angles, cart_coords, batch, sigma):
        edge_index, distance_vectors, _ = get_pbc_cutoff_graphs(cart_coords, lengths, angles, num_atoms, self.cutoff, self.max_num_neighbors)
        edge_features = self.backbone(atom_types, edge_index, distance_vectors)
        num_multi_edge_per_graph = scatter(torch.ones(size=(edge_index.shape[1],), device=edge_index.device).long(), batch[edge_index[0]], reduce='sum')
        latents_per_multi_edge = latents.repeat_interleave(num_multi_edge_per_graph, dim=0)
        edge_features = torch.cat((edge_features, latents_per_multi_edge), dim=1)

        j, i = edge_index
        no_iden_mask = (i != j)
        j, i, edge_features, distance_vectors = j[no_iden_mask], i[no_iden_mask], edge_features[no_iden_mask], distance_vectors[no_iden_mask]
        dists_per_multi_edge = torch.linalg.norm(distance_vectors, dim=-1, keepdim=True)
        scores_per_multi_edge = self.fc_score(edge_features) * distance_vectors / dists_per_multi_edge

        num_multi_edges = len(i)
        new_edge_start_mask = torch.logical_or(i[:-1] != i[1:], j[:-1] != j[1:])
        new_edge_start_id = torch.nonzero(new_edge_start_mask).view(-1) + 1
        num_multi_edges_per_edge = torch.cat([new_edge_start_id[0:1], new_edge_start_id[1:] - new_edge_start_id[:-1], num_multi_edges - new_edge_start_id[-1:]])
        multi_edge_to_edge_idx = torch.repeat_interleave(torch.arange(len(num_multi_edges_per_edge), device=num_multi_edges_per_edge.device), num_multi_edges_per_edge)
        scores_per_edge = scatter(scores_per_multi_edge, multi_edge_to_edge_idx, dim=0, reduce='mean')
        unique_edge_receiver_index = scatter(i, multi_edge_to_edge_idx, dim=0, reduce='mean').long()
        scores_per_node_pos = scatter(scores_per_edge, unique_edge_receiver_index, dim=0, dim_size=len(batch), reduce='sum')

        score_norm = self.get_score_norm(sigma)
        return scores_per_node_pos / score_norm
//  [python]
def compute_normals(inputs: torch.Tensor,
                    outputs: torch.Tensor,
                    eps: float = 1e-4,
                    grad_min: float = -1e4,
                    grad_max: float = 1e4,
                    retain_graph: bool = True):
    """
    Args:
        inputs (B, N, d):
        outputs (B, N):

    Returns:

    """
    assert inputs.ndim == 3 and outputs.ndim == 2, 'input shape expeted to be [B, N, d], and output to be [B, N]'
    d_points = torch.ones_like(outputs, requires_grad=False, device=outputs.device)
    points_grad = torch.autograd.grad(
        outputs=outputs,
        inputs=inputs,
        grad_outputs=d_points,
        create_graph=retain_graph,
        retain_graph=retain_graph,
        only_inputs=True)[0][:, :, -3:]

    points_grad = torch.clamp(points_grad, grad_min, grad_max)
    norm = torch.norm(points_grad, dim=-1).unsqueeze(-1).detach()
    point_normals = points_grad / (eps + norm)

    return point_normals
//  [python]
def compute_normals(inputs: torch.Tensor,
                    outputs: torch.Tensor,
                    eps: float = 1e-4,
                    grad_min: float = -1e4,
                    grad_max: float = 1e4,
                    retain_graph: bool = True):
    """
    Args:
        inputs (B, N, d):
        outputs (B, N):

    Returns:

    """
    assert inputs.ndim == 3 and outputs.ndim == 2, 'input shape expeted to be [B, N, d], and output to be [B, N]'
    d_points = torch.ones_like(outputs, requires_grad=False, device=outputs.device)
    points_grad = torch.autograd.grad(
        outputs=outputs,
        inputs=inputs,
        grad_outputs=d_points,
        create_graph=retain_graph,
        retain_graph=retain_graph,
        only_inputs=True)[0][:, :, -3:]

    points_grad = torch.clamp(points_grad, grad_min, grad_max)
    norm = torch.linalg.norm(points_grad, dim=-1).unsqueeze(-1).detach()
    point_normals = points_grad / (eps + norm)

    return point_normals
//  [python]
def optimize(moduleList : list,learning_rate : float = 2 ** -7,signSGDfactor : float = 1):
    """
    Interpolates between signSGD and infinity norm normalization.

    signSGDfactor - The interpolation factor for signSGD. 0 is pure infinity norm normalization, while 1 is pure signSGD.
    """
    assert (0 <= signSGDfactor <= 1)
    assert (0 <= learning_rate)
    with torch.no_grad():
        for module in moduleList:
            infNorm = torch.scalar_tensor(0)
            paramGroup = [x for x in module.parameters() if x.grad is not None]
            for x in paramGroup: infNorm = torch.max(infNorm,torch.norm(x.grad.flatten(),float('inf')))
            for x in paramGroup: # The scalars operations are grouped for efficiency.
                if torch.is_nonzero(infNorm):
                    if 0 < signSGDfactor: x -= learning_rate * signSGDfactor * torch.sign(x.grad)
                    if signSGDfactor < 1: x -= learning_rate * (1 - signSGDfactor) / infNorm * x.grad
                    x.grad = None
//  [python]
def optimize(moduleList : list,learning_rate : float = 2 ** -7,signSGDfactor : float = 1):
    """
    Interpolates between signSGD and infinity norm normalization.

    signSGDfactor - The interpolation factor for signSGD. 0 is pure infinity norm normalization, while 1 is pure signSGD.
    """
    assert (0 <= signSGDfactor <= 1)
    assert (0 <= learning_rate)
    with torch.no_grad():
        for module in moduleList:
            infNorm = torch.scalar_tensor(0)
            paramGroup = [x for x in module.parameters() if x.grad is not None]
            for x in paramGroup: infNorm = torch.max(infNorm,torch.linalg.norm(x.grad.flatten(),ord=float('inf')))
            for x in paramGroup: # The scalars operations are grouped for efficiency.
                if torch.is_nonzero(infNorm):
                    if 0 < signSGDfactor: x -= learning_rate * signSGDfactor * torch.sign(x.grad)
                    if signSGDfactor < 1: x -= learning_rate * (1 - signSGDfactor) / infNorm * x.grad
                    x.grad = None
//  [python]
def to_point_direction(self, normalize_moment=True):
        """
        Convert to point direction representation <O, D>.

        Returns:
            rays: (..., 6).
        """
        if self._is_plucker:
            direction = torch.nn.functional.normalize(self.rays[..., :3], dim=-1)
            moment = self.rays[..., 3:]
            if normalize_moment:
                c = torch.norm(direction, dim=-1, keepdim=True)
                moment = moment / c
            points = torch.cross(direction, moment, dim=-1)
            return Rays(
                rays=torch.cat((points, direction), dim=-1),
                is_plucker=False,
                ndc_coordinates=self.ndc_coordinates,
            )
        else:
            return self
//  [python]
def to_point_direction(self, normalize_moment=True):
        """
        Convert to point direction representation <O, D>.

        Returns:
            rays: (..., 6).
        """
        if self._is_plucker:
            direction = torch.nn.functional.normalize(self.rays[..., :3], dim=-1)
            moment = self.rays[..., 3:]
            if normalize_moment:
                c = torch.linalg.norm(direction, dim=-1, keepdim=True)
                moment = moment / c
            points = torch.cross(direction, moment, dim=-1)
            return Rays(
                rays=torch.cat((points, direction), dim=-1),
                is_plucker=False,
                ndc_coordinates=self.ndc_coordinates,
            )
        else:
            return self
//  [python]
def _compute_local_normal_vector(self, end_corner, start_corner):
        conect_vector = end_corner - start_corner
        side_length = torch.norm(conect_vector, dim=0)
        normal = torch.index_select(conect_vector, 0, torch.LongTensor([1, 0]))
        normal /= side_length
        normal[1] *= -1
        return normal
//  [python]
def _compute_local_normal_vector(self, end_corner, start_corner):
        conect_vector = end_corner - start_corner
        side_length = torch.linalg.norm(conect_vector, dim=0)
        normal = torch.index_select(conect_vector, 0, torch.LongTensor([1, 0]))
        normal /= side_length
        normal[1] *= -1
        return normal
//  [python]
def _distribute_line_to_boundary(self, points, index, line_points,
                                     corners, current_length):
        corner_index = 0
        side_length = torch.norm(corners[1] - corners[0])
        while index < len(line_points):
            if line_points[index] <= current_length + side_length:
                point = self._translate_point_to_bondary(index, line_points,
                                                         corners, current_length,
                                                         corner_index, side_length)
                points[index] = point
                index += 1
            else:
                corner_index += 1
                current_length += side_length
                if corner_index >= len(corners) - 1:
                    break
                side_length = torch.norm(corners[corner_index+1]
                                             - corners[corner_index])
        return points, index, current_length
//  [python]
def _distribute_line_to_boundary(self, points, index, line_points,
                                     corners, current_length):
        corner_index = 0
        side_length = torch.linalg.norm(corners[1] - corners[0])
        while index < len(line_points):
            if line_points[index] <= current_length + side_length:
                point = self._translate_point_to_bondary(index, line_points,
                                                         corners, current_length,
                                                         corner_index, side_length)
                points[index] = point
                index += 1
            else:
                corner_index += 1
                current_length += side_length
                if corner_index >= len(corners) - 1:
                    break
                side_length = torch.linalg.norm(corners[corner_index+1]
                                             - corners[corner_index])
        return points, index, current_length
//  [python]
def to_stiefel_polar(theta, dim:int, rank:int):
    r'''map real vector to a Stiefel manifold via polar decomposition

    [wiki-link](https://en.wikipedia.org/wiki/Stiefel_manifold)

    $$ \left\{ x\in\mathbb{K}^{m\times n}:x^\dagger x=I_n \right\}$$

    where $\mathbb{K}=\mathbb{R}$ or $\mathbb{C}$. Matrix square root is used so the backward might be
    several times slower than the forward. QR decomposition

    Parameters:
        theta (np.ndarray,torch.Tensor): if `ndim>1`, then the last dimension will be expanded to the matrix
                and the rest dimensions will be batch dimensions.
        dim (int): dimension of the matrix.
        rank (int): rank of the matrix.

    Returns:
        ret (np.ndarray,torch.Tensor): array of shape `theta.shape[:-1]+(dim,rank)`
    '''
    assert rank<=dim
    shape = theta.shape
    if shape[-1]==dim*rank: #real
        is_real = True
    else:
        assert shape[-1]==2*dim*rank
        is_real = False
    theta = theta.reshape(-1,shape[-1])
    if isinstance(theta, torch.Tensor):
        if is_real:
            mat = theta.reshape(-1, dim, rank)
        else:
            tmp0 = theta.reshape(-1,2,dim,rank)
            mat = torch.complex(tmp0[:,0], tmp0[:,1])
        if rank==1:
            ret = mat / torch.norm(mat, dim=1, keepdim=True)
        else:
            tmp0 = torch.linalg.inv(numqi._torch_op.PSDMatrixSqrtm.apply(mat.transpose(1,2).conj() @ mat))
            ret = mat @ tmp0
    else: #numpy
        if is_real:
            mat = theta.reshape(-1, dim, rank)
        else:
            tmp0 = theta.reshape(-1,2,dim,rank)
            mat = tmp0[:,0] + 1j*tmp0[:,1]
        if rank==1:
            ret = mat / np.linalg.norm(mat, axis=1, keepdims=True)
        else:
            # scipy.linalg.sqrtm is slow, so we use np.linalg.eigh here
            EVL,EVC = np.linalg.eigh(mat.transpose(0,2,1).conj() @ mat)
            tmp0 = (EVC*np.sqrt(1/EVL).reshape(-1,1,rank)) @ EVC.transpose(0,2,1).conj()
            ret = mat @ tmp0
    ret = ret.reshape(*shape[:-1], dim, rank)
    return ret
//  [python]
def to_stiefel_polar(theta, dim:int, rank:int):
    r'''map real vector to a Stiefel manifold via polar decomposition

    [wiki-link](https://en.wikipedia.org/wiki/Stiefel_manifold)

    $$ \left\{ x\in\mathbb{K}^{m\times n}:x^\dagger x=I_n \right\}$$

    where $\mathbb{K}=\mathbb{R}$ or $\mathbb{C}$. Matrix square root is used so the backward might be
    several times slower than the forward. QR decomposition

    Parameters:
        theta (np.ndarray,torch.Tensor): if `ndim>1`, then the last dimension will be expanded to the matrix
                and the rest dimensions will be batch dimensions.
        dim (int): dimension of the matrix.
        rank (int): rank of the matrix.

    Returns:
        ret (np.ndarray,torch.Tensor): array of shape `theta.shape[:-1]+(dim,rank)`
    '''
    assert rank<=dim
    shape = theta.shape
    if shape[-1]==dim*rank: #real
        is_real = True
    else:
        assert shape[-1]==2*dim*rank
        is_real = False
    theta = theta.reshape(-1,shape[-1])
    if isinstance(theta, torch.Tensor):
        if is_real:
            mat = theta.reshape(-1, dim, rank)
        else:
            tmp0 = theta.reshape(-1,2,dim,rank)
            mat = torch.complex(tmp0[:,0], tmp0[:,1])
        if rank==1:
            ret = mat / torch.linalg.norm(mat, axis=1, keepdims=True)
        else:
            tmp0 = torch.linalg.inv(numqi._torch_op.PSDMatrixSqrtm.apply(mat.transpose(1,2).conj() @ mat))
            ret = mat @ tmp0
    else: #numpy
        if is_real:
            mat = theta.reshape(-1, dim, rank)
        else:
            tmp0 = theta.reshape(-1,2,dim,rank)
            mat = tmp0[:,0] + 1j*tmp0[:,1]
        if rank==1:
            ret = mat / np.linalg.norm(mat, axis=1, keepdims=True)
        else:
            # scipy.linalg.sqrtm is slow, so we use np.linalg.eigh here
            EVL,EVC = np.linalg.eigh(mat.transpose(0,2,1).conj() @ mat)
            tmp0 = (EVC*np.sqrt(1/EVL).reshape(-1,1,rank)) @ EVC.transpose(0,2,1).conj()
            ret = mat @ tmp0
    ret = ret.reshape(*shape[:-1], dim, rank)
    return ret
//  [python]
def estimate_M(
		Ys, Xs, M, sigma_yxs, betas, M_constraint, context,
		n_epochs=10000, tol=None, verbose=True,
):
	"""
	M is shared across all replicates
	min || Y - X MT ||_2^2 / (2 σ_yx^2)
	s.t. || Mk ||_p = 1
	grad = (M XT X - YT X) / (σ_yx^2)

	Each replicate may have a slightly different M
	min || Y - X MT ||_2^2 / (2 σ_yx^2) + || M - M_base ||_2^2 λ_M / 2
	s.t. || Mk ||_p = 1
	grad = ( M XT X - YT X ) / ( σ_yx^2 ) + λ_M ( M - M_base )
	"""
	if tol is None:
		tol = 1e-5
	tol /= len(M)
	K = M.shape[1]
	quad = torch.zeros([K, K], **context)
	linear = torch.zeros_like(M)
	constant = (np.array([torch.norm(Y).item()**2 for Y in Ys]) * betas / (sigma_yxs**2)).sum()
	for Y, X, sigma_yx, beta in zip(Ys, Xs, sigma_yxs, betas):
		quad.addmm_(X.T, X, alpha=beta / (sigma_yx ** 2))
		linear.addmm_(Y.T.to(X.device), X, alpha=beta / (sigma_yx ** 2))

	loss_prev, loss = np.inf, np.nan
	pbar = trange(1, n_epochs+1, leave=True, disable=not verbose, desc='Updating M', miniters=1000)

	def calc_func(M):
		# f = (M @ quad).ravel() @ M.ravel() / 2
		# f -= linear.ravel() @ M.ravel()
		f = torch.addmm(linear, M, quad, alpha=.5, beta=-1).ravel() @ M.ravel() + constant / 2
		return f
	def calc_grad(M):
		g = torch.addmm(linear, M, quad, alpha=1, beta=-1)
		return g
	step_size = 1 / torch.eigvalsh(quad).max().item()
	optimizer = NesterovGD(M, step_size)
	valid_freq = 100
	M_prev = M.clone()
	func, func_prev, df, dM = np.nan, np.inf, np.nan, np.inf
	for i_epoch in pbar:
		grad = calc_grad(M)
		optimizer.step(grad)
		project_M(M, M_constraint)
		if i_epoch % valid_freq == 0 or i_epoch == n_epochs:
			func = calc_func(M).item()
			df = func_prev - func
			func_prev = func
			dM = M_prev.sub_(M).abs_().max().item()
			M_prev = M.clone()
			do_stop = dM < tol * np.sqrt(valid_freq) and i_epoch > 5
			assert np.isfinite(func)
			pbar.set_description(
				f'Updating M: loss = {func:.1e}, '
				f'%δloss = {df / func:.1e}, '
				f'δM = {dM:.1e}'
			)
			if do_stop:
				break
//  [python]
def estimate_M(
		Ys, Xs, M, sigma_yxs, betas, M_constraint, context,
		n_epochs=10000, tol=None, verbose=True,
):
	"""
	M is shared across all replicates
	min || Y - X MT ||_2^2 / (2 σ_yx^2)
	s.t. || Mk ||_p = 1
	grad = (M XT X - YT X) / (σ_yx^2)

	Each replicate may have a slightly different M
	min || Y - X MT ||_2^2 / (2 σ_yx^2) + || M - M_base ||_2^2 λ_M / 2
	s.t. || Mk ||_p = 1
	grad = ( M XT X - YT X ) / ( σ_yx^2 ) + λ_M ( M - M_base )
	"""
	if tol is None:
		tol = 1e-5
	tol /= len(M)
	K = M.shape[1]
	quad = torch.zeros([K, K], **context)
	linear = torch.zeros_like(M)
	constant = (np.array([torch.linalg.norm(Y, ord='fro').item()**2 for Y in Ys]) * betas / (sigma_yxs**2)).sum()
	for Y, X, sigma_yx, beta in zip(Ys, Xs, sigma_yxs, betas):
		quad.addmm_(X.T, X, alpha=beta / (sigma_yx ** 2))
		linear.addmm_(Y.T.to(X.device), X, alpha=beta / (sigma_yx ** 2))

	loss_prev, loss = np.inf, np.nan
	pbar = trange(1, n_epochs+1, leave=True, disable=not verbose, desc='Updating M', miniters=1000)

	def calc_func(M):
		# f = (M @ quad).ravel() @ M.ravel() / 2
		# f -= linear.ravel() @ M.ravel()
		f = torch.addmm(linear, M, quad, alpha=.5, beta=-1).ravel() @ M.ravel() + constant / 2
		return f
	def calc_grad(M):
		g = torch.addmm(linear, M, quad, alpha=1, beta=-1)
		return g
	step_size = 1 / torch.linalg.eigvalsh(quad).max().item()
	optimizer = NesterovGD(M, step_size)
	valid_freq = 100
	M_prev = M.clone()
	func, func_prev, df, dM = np.nan, np.inf, np.nan, np.inf
	for i_epoch in pbar:
		grad = calc_grad(M)
		optimizer.step(grad)
		project_M(M, M_constraint)
		if i_epoch % valid_freq == 0 or i_epoch == n_epochs:
			func = calc_func(M).item()
			df = func_prev - func
			func_prev = func
			dM = M_prev.sub_(M).abs_().max().item()
			M_prev = M.clone()
			do_stop = dM < tol * np.sqrt(valid_freq) and i_epoch > 5
			assert np.isfinite(func)
			pbar.set_description(
				f'Updating M: loss = {func:.1e}, '
				f'%δloss = {df / func:.1e}, '
				f'δM = {dM:.1e}'
			)
			if do_stop:
				break
//  [python]
def dq_to_x(
        self,
        coords_x: torch.Tensor,
        dq: torch.Tensor,
        constr_tol: float = 0.1,
        enforce_constr: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor, bool]:
        coords_x, coords_q, q_converged = _dq_to_x_cached(coords_x, dq, self)

        if self.constr is None:
            return coords_x, coords_q, q_converged

        constr_delta = self.compute_constr_delta(coords_x)

        if enforce_constr and torch.norm(constr_delta) >= constr_tol:
            return self._apply_constraints(coords_x, constr_tol)

        return coords_x, coords_q, q_converged
//  [python]
def dq_to_x(
        self,
        coords_x: torch.Tensor,
        dq: torch.Tensor,
        constr_tol: float = 0.1,
        enforce_constr: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor, bool]:
        coords_x, coords_q, q_converged = _dq_to_x_cached(coords_x, dq, self)

        if self.constr is None:
            return coords_x, coords_q, q_converged

        constr_delta = self.compute_constr_delta(coords_x)

        if enforce_constr and torch.linalg.norm(constr_delta) >= constr_tol:
            return self._apply_constraints(coords_x, constr_tol)

        return coords_x, coords_q, q_converged
//  [python]
def _dq_to_x_cached(
    coords_x: torch.Tensor,
    dq: torch.Tensor,
    ic: typing.Union["RIC", "DLC"],
) -> tuple[torch.Tensor, torch.Tensor, bool]:
    """A cached version of the ``_dq_to_x`` function.

    While a bit of a hack, this can significantly speed up ops like the brent
    line search.
    """
    if ic._dq_cache is not None:
        (
            cached_coords_x,
            cached_dq,
            cached_coords_x_new,
            cached_coords_q_new,
            cached_q_converged,
        ) = ic._dq_cache

        if (
            torch.norm(cached_coords_x - coords_x) < 1e-10
            and torch.norm(cached_dq - dq) < 1e-10
        ):
            return cached_coords_x_new, cached_coords_q_new, cached_q_converged

    coords_x_new, coords_q_new, q_converged = _dq_to_x(coords_x, dq, ic)
    ic._dq_cache = (coords_x, dq, coords_x_new, coords_q_new, q_converged)

    return coords_x_new, coords_q_new, q_converged