jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def gen_rays_at(self, img_idx, resolution_level=1):
        """
        Generate rays at world space from one camera.
        """

        src_idx = self.src_idx[img_idx]
        src_idx = src_idx[:9]
        # src_idx = src_idx[-9:]
        idx_list = torch.cat([torch.tensor(img_idx).unsqueeze(0), src_idx], dim=0)

        poses_pair = self.pose_all[idx_list]  # [store R^-1 and C]
        intrinsics_pair = self.intrinsics_all[idx_list]
        intrinsics_inv_pair = self.intrinsics_all_inv[idx_list]
        images_gray_pair = self.images_gray[idx_list]

        l = resolution_level
        tx = torch.linspace(0, self.W - 1, self.W // l)
        ty = torch.linspace(0, self.H - 1, self.H // l)
        pixels_x, pixels_y = torch.meshgrid(tx, ty)

        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1) # W, H, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, None, :3, :3], p[:, :, :, None]).squeeze()  # W, H, 3
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)  # W, H, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, None, :3, :3], rays_v[:, :, :, None]).squeeze()  # W, H, 3
        rays_o = self.pose_all[img_idx, None, None, :3, 3].expand(rays_v.shape)  # W, H, 3
        return rays_o.transpose(0, 1), rays_v.transpose(0, 1), intrinsics_pair, intrinsics_inv_pair, poses_pair, images_gray_pair
//  [python]
def move_until_success(
        self, position, orientation, time_to_go, max_attempts=5, success_dist=0.05
    ):
        states = []
        for _ in range(max_attempts):
            joint_pos = self.ik(position, orientation)
            states += self.move_to_joint_positions(joint_pos, time_to_go=time_to_go)
            curr_ee_pos, curr_ee_ori = self.get_ee_pose()

            xyz_diff = t.norm(curr_ee_pos - position)
            ori_diff = (
                R.from_quat(curr_ee_ori).inv() * R.from_quat(orientation)
            ).magnitude()
            log.info(f"Dist to goal: xyz diff {xyz_diff}, ori diff {ori_diff}")

            if (
                states
                and states[-1].prev_command_successful
                and xyz_diff < success_dist
                and ori_diff < 0.2
            ):
                break
        return states
//  [python]
def move_until_success(
        self, position, orientation, time_to_go, max_attempts=5, success_dist=0.05
    ):
        states = []
        for _ in range(max_attempts):
            joint_pos = self.ik(position, orientation)
            states += self.move_to_joint_positions(joint_pos, time_to_go=time_to_go)
            curr_ee_pos, curr_ee_ori = self.get_ee_pose()

            xyz_diff = torch.linalg.norm(curr_ee_pos - position)
            ori_diff = (
                R.from_quat(curr_ee_ori).inv() * R.from_quat(orientation)
            ).magnitude()
            log.info(f"Dist to goal: xyz diff {xyz_diff}, ori diff {ori_diff}")

            if (
                states
                and states[-1].prev_command_successful
                and xyz_diff < success_dist
                and ori_diff < 0.2
            ):
                break
        return states
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """

        src_idx = self.src_idx[img_idx]
        src_idx = src_idx[:9]
        # src_idx = src_idx[-9:]
        idx_list = torch.cat([img_idx.clone().detach().unsqueeze(0), src_idx], dim=0).cuda()

        poses_pair = self.pose_all[idx_list]  # [store R^-1 and C]
        intrinsics_pair = self.intrinsics_all[idx_list]
        intrinsics_inv_pair = self.intrinsics_all_inv[idx_list]
        images_gray_pair = self.images_gray[idx_list]

        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size])
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])
        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze() # batch_size, 3
        rays_v = p / torch.norm(p, p=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape) # batch_size, 3
        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask[:, :1]], dim=-1).cuda(), intrinsics_pair, intrinsics_inv_pair, poses_pair, images_gray_pair
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """

        src_idx = self.src_idx[img_idx]
        src_idx = src_idx[:9]
        # src_idx = src_idx[-9:]
        idx_list = torch.cat([img_idx.clone().detach().unsqueeze(0), src_idx], dim=0).cuda()

        poses_pair = self.pose_all[idx_list]  # [store R^-1 and C]
        intrinsics_pair = self.intrinsics_all[idx_list]
        intrinsics_inv_pair = self.intrinsics_all_inv[idx_list]
        images_gray_pair = self.images_gray[idx_list]

        pixels_x = torch.randint(low=0, high=self.W, size=[batch_size])
        pixels_y = torch.randint(low=0, high=self.H, size=[batch_size])
        color = self.images[img_idx][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx][(pixels_y, pixels_x)]      # batch_size, 3
        p = torch.stack([pixels_x, pixels_y, torch.ones_like(pixels_y)], dim=-1).float()  # batch_size, 3
        p = torch.matmul(self.intrinsics_all_inv[img_idx, None, :3, :3], p[:, :, None]).squeeze() # batch_size, 3
        rays_v = p / torch.linalg.norm(p, ord=2, dim=-1, keepdim=True)    # batch_size, 3
        rays_v = torch.matmul(self.pose_all[img_idx, None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[img_idx, None, :3, 3].expand(rays_v.shape) # batch_size, 3
        return torch.cat([rays_o.cpu(), rays_v.cpu(), color, mask[:, :1]], dim=-1).cuda(), intrinsics_pair, intrinsics_inv_pair, poses_pair, images_gray_pair
//  [python]
def get_edge_vectors_and_lengths(
    positions: torch.Tensor,  # [n_nodes, 3]
    edge_index: torch.Tensor,  # [2, n_edges]
    shifts: torch.Tensor,  # [n_edges, 3]
    normalize: bool = False,
    eps: float = 1e-9,
) -> Tuple[torch.Tensor, torch.Tensor]:
    sender = edge_index[0]
    receiver = edge_index[1]
    vectors = positions[receiver] - positions[sender] + shifts  # [n_edges, 3]
    lengths = torch.norm(vectors, dim=-1, keepdim=True)  # [n_edges, 1]
    if normalize:
        vectors_normed = vectors / (lengths + eps)
        return vectors_normed, lengths

    return vectors, lengths
//  [python]
def get_edge_vectors_and_lengths(
    positions: torch.Tensor,  # [n_nodes, 3]
    edge_index: torch.Tensor,  # [2, n_edges]
    shifts: torch.Tensor,  # [n_edges, 3]
    normalize: bool = False,
    eps: float = 1e-9,
) -> Tuple[torch.Tensor, torch.Tensor]:
    sender = edge_index[0]
    receiver = edge_index[1]
    vectors = positions[receiver] - positions[sender] + shifts  # [n_edges, 3]
    lengths = torch.linalg.norm(vectors, dim=-1, keepdim=True)  # [n_edges, 1]
    if normalize:
        vectors_normed = vectors / (lengths + eps)
        return vectors_normed, lengths

    return vectors, lengths
//  [python]
def compute_scale_loss(self, translation_scale, expec_rt_gt):
        scale_gt = torch.norm(expec_rt_gt[0][...,:3,3])

        scale_threshold = self.loss_config['max_scale_loss']

        power = 2
        loss_scale = torch.pow(translation_scale-scale_gt,power).mean()
        loss_scale_reported = min(loss_scale, torch.tensor(scale_threshold).to(device=translation_scale.device))
        if loss_scale_reported == torch.tensor(scale_threshold).to(device=translation_scale.device):
            # making backprop legit
            loss_scale *= 0

        return loss_scale, loss_scale_reported
//  [python]
def compute_scale_loss(self, translation_scale, expec_rt_gt):
        scale_gt = torch.linalg.norm(expec_rt_gt[0][...,:3,3])

        scale_threshold = self.loss_config['max_scale_loss']

        power = 2
        loss_scale = torch.pow(translation_scale-scale_gt,power).mean()
        loss_scale_reported = min(loss_scale, torch.tensor(scale_threshold).to(device=translation_scale.device))
        if loss_scale_reported == torch.tensor(scale_threshold).to(device=translation_scale.device):
            # making backprop legit
            loss_scale *= 0

        return loss_scale, loss_scale_reported
//  [python]
def rot_param_rot_mat(rot):
    '''

    :param rot: N x 6
    :return:
    '''

    e1 = rot[:,:3]
    e2 = rot[:,3:]

    e1d = e1/torch.norm(e1,dim=1,keepdim=True)
    e3d = torch.cross(e1d, e2, dim=1)/torch.norm(e2,dim=1,keepdim=True)
    e2d = torch.cross(e3d, e1d)

    R = torch.stack([e1d, e2d, e3d],dim=2) # N x 3 x 3

    return R
//  [python]
def rot_param_rot_mat(rot):
    '''

    :param rot: N x 6
    :return:
    '''

    e1 = rot[:,:3]
    e2 = rot[:,3:]

    e1d = e1/torch.linalg.norm(e1,dim=1,keepdim=True)
    e3d = torch.cross(e1d, e2, dim=1)/torch.linalg.norm(e2,dim=1,keepdim=True)
    e2d = torch.cross(e3d, e1d)

    R = torch.stack([e1d, e2d, e3d],dim=2) # N x 3 x 3

    return R
//  [python]
def _compute_fine_loss_l2_std(self, expec_f, expec_f_gt):
        """
        Args:
            expec_f (torch.Tensor): [M, 3] <x, y, std>
            expec_f_gt (torch.Tensor): [M, 2] <x, y>
        """
        #import pdb; pdb.set_trace()
        # correct_mask tells you which pair to compute fine-loss
        correct_mask = torch.norm(expec_f_gt, p=float('inf'), dim=1) < self.correct_thr

        # use std as weight that measures uncertainty
        std = expec_f[:, 2]
        inverse_std = 1. / torch.clamp(std, min=1e-10)
        weight = (inverse_std / torch.mean(inverse_std)).detach()  # avoid minizing loss through increase std

        # corner case: no correct coarse match found
        if not correct_mask.any():
            if self.training:  # this seldomly happen during training, since we pad prediction with gt
                               # sometimes there is not coarse-level gt at all.
                logger.warning("assign a false supervision to avoid ddp deadlock")
                correct_mask[0] = True
                weight[0] = 0.
            else:
                return None

        # l2 loss with std
        offset_l2 = ((expec_f_gt[correct_mask] - expec_f[correct_mask, :2]) ** 2).sum(-1)
        loss = (offset_l2 * weight[correct_mask]).mean()

        return loss
//  [python]
def _compute_fine_loss_l2_std(self, expec_f, expec_f_gt):
        """
        Args:
            expec_f (torch.Tensor): [M, 3] <x, y, std>
            expec_f_gt (torch.Tensor): [M, 2] <x, y>
        """
        #import pdb; pdb.set_trace()
        # correct_mask tells you which pair to compute fine-loss
        correct_mask = torch.linalg.norm(expec_f_gt, ord=float('inf'), dim=1) < self.correct_thr

        # use std as weight that measures uncertainty
        std = expec_f[:, 2]
        inverse_std = 1. / torch.clamp(std, min=1e-10)
        weight = (inverse_std / torch.mean(inverse_std)).detach()  # avoid minizing loss through increase std

        # corner case: no correct coarse match found
        if not correct_mask.any():
            if self.training:  # this seldomly happen during training, since we pad prediction with gt
                               # sometimes there is not coarse-level gt at all.
                logger.warning("assign a false supervision to avoid ddp deadlock")
                correct_mask[0] = True
                weight[0] = 0.
            else:
                return None

        # l2 loss with std
        offset_l2 = ((expec_f_gt[correct_mask] - expec_f[correct_mask, :2]) ** 2).sum(-1)
        loss = (offset_l2 * weight[correct_mask]).mean()

        return loss
//  [python]
def get_plane_transform(up, ground_plane=None, xyz_orig=None):
    """
    get R, t rigid transform from plane and desired origin
    :param up (3,) up vector of coordinate frame
    :param ground_plane (4) (a, b, c, d) where a,b,c is the normal
    :param xyz_orig (3) desired origin
    """
    R = torch.eye(3)
    t = torch.zeros(3)
    if ground_plane is None:
        return R, t

    # compute transform between world up vector and passed in floor
    ground_plane = torch.as_tensor(ground_plane)
    ground_plane = torch.sign(ground_plane[3]) * ground_plane

    normal = ground_plane[:3]
    normal = normal / torch.norm(normal)
    v = torch.linalg.cross(up, normal)
    ang_sin = torch.norm(v)
    ang_cos = up.dot(normal)
    skew_v = torch.as_tensor([[0.0, -v[2], v[1]], [v[2], 0.0, -v[0]], [-v[1], v[0], 0.0]])
    R = torch.eye(3) + skew_v + (skew_v @ skew_v) * ((1.0 - ang_cos) / (ang_sin**2))

    # project origin onto plane
    if xyz_orig is None:
        xyz_orig = torch.zeros(3)
    t, _ = compute_plane_intersection(xyz_orig, -normal, ground_plane)

    return R, t
//  [python]
def get_plane_transform(up, ground_plane=None, xyz_orig=None):
    """
    get R, t rigid transform from plane and desired origin
    :param up (3,) up vector of coordinate frame
    :param ground_plane (4) (a, b, c, d) where a,b,c is the normal
    :param xyz_orig (3) desired origin
    """
    R = torch.eye(3)
    t = torch.zeros(3)
    if ground_plane is None:
        return R, t

    # compute transform between world up vector and passed in floor
    ground_plane = torch.as_tensor(ground_plane)
    ground_plane = torch.sign(ground_plane[3]) * ground_plane

    normal = ground_plane[:3]
    normal = normal / torch.linalg.norm(normal)
    v = torch.linalg.cross(up, normal)
    ang_sin = torch.linalg.norm(v)
    ang_cos = up.dot(normal)
    skew_v = torch.as_tensor([[0.0, -v[2], v[1]], [v[2], 0.0, -v[0]], [-v[1], v[0], 0.0]])
    R = torch.eye(3) + skew_v + (skew_v @ skew_v) * ((1.0 - ang_cos) / (ang_sin**2))

    # project origin onto plane
    if xyz_orig is None:
        xyz_orig = torch.zeros(3)
    t, _ = compute_plane_intersection(xyz_orig, -normal, ground_plane)

    return R, t
//  [python]
def _guess_hess_distance(
    coords: torch.Tensor, atomic_nums: torch.Tensor, idxs: torch.Tensor
) -> list[float]:
    dist = torch.norm(coords[idxs[:, 0]] - coords[idxs[:, 1]], dim=1)
    hess = []

    for (idx_a, idx_b), r in zip(idxs, dist, strict=True):
        a = 1.734

        z_1 = min(atomic_nums[idx_a], atomic_nums[idx_b])
        z_2 = max(atomic_nums[idx_a], atomic_nums[idx_b])

        if z_1 < 3:
            if z_2 < 3:
                b = -0.244
            elif z_2 < 11:
                b = 0.352
            else:
                b = 0.660
        elif z_1 < 11:
            if z_2 < 11:
                b = 1.085
            else:
                b = 1.522
        else:
            b = 2.068

        hess.append((a / (r - b) ** 3).detach().cpu().item())

    return hess
//  [python]
def _guess_hess_distance(
    coords: torch.Tensor, atomic_nums: torch.Tensor, idxs: torch.Tensor
) -> list[float]:
    dist = torch.linalg.norm(coords[idxs[:, 0]] - coords[idxs[:, 1]], dim=1)
    hess = []

    for (idx_a, idx_b), r in zip(idxs, dist, strict=True):
        a = 1.734

        z_1 = min(atomic_nums[idx_a], atomic_nums[idx_b])
        z_2 = max(atomic_nums[idx_a], atomic_nums[idx_b])

        if z_1 < 3:
            if z_2 < 3:
                b = -0.244
            elif z_2 < 11:
                b = 0.352
            else:
                b = 0.660
        elif z_1 < 11:
            if z_2 < 11:
                b = 1.085
            else:
                b = 1.522
        else:
            b = 2.068

        hess.append((a / (r - b) ** 3).detach().cpu().item())

    return hess
//  [python]
def comp_point_point_error(self, Xt, Yt):
        closest_idx = self.comp_closest_pts_idx_with_split(Xt, Yt)
        pt_pt_vec = Xt - Yt[:, closest_idx]  # (3, S) - (3, S) -> (3, S)
        pt_pt_dist = torch.norm(pt_pt_vec, dim=0)
        eng = torch.mean(pt_pt_dist)
        return eng
//  [python]
def comp_point_point_error(self, Xt, Yt):
        closest_idx = self.comp_closest_pts_idx_with_split(Xt, Yt)
        pt_pt_vec = Xt - Yt[:, closest_idx]  # (3, S) - (3, S) -> (3, S)
        pt_pt_dist = torch.linalg.norm(pt_pt_vec, dim=0)
        eng = torch.mean(pt_pt_dist)
        return eng
//  [python]
def comp_closest_pts_idx_with_split(self, pts_src, pts_des):
        """
        :param pts_src:     (3, S)
        :param pts_des:     (3, D)
        :param num_split:
        :return:
        """
        pts_src_list = torch.split(pts_src, 500000, dim=1)
        idx_list = []
        for pts_src_sec in pts_src_list:
            diff = pts_src_sec[:, :, np.newaxis] - pts_des[:, np.newaxis, :]  # (3, S, 1) - (3, 1, D) -> (3, S, D)
            dist = torch.norm(diff, dim=0)  # (S, D)
            closest_idx = torch.argmin(dist, dim=1)  # (S,)
            idx_list.append(closest_idx)
        closest_idx = torch.cat(idx_list)
        return closest_idx
//  [python]
def comp_closest_pts_idx_with_split(self, pts_src, pts_des):
        """
        :param pts_src:     (3, S)
        :param pts_des:     (3, D)
        :param num_split:
        :return:
        """
        pts_src_list = torch.split(pts_src, 500000, dim=1)
        idx_list = []
        for pts_src_sec in pts_src_list:
            diff = pts_src_sec[:, :, np.newaxis] - pts_des[:, np.newaxis, :]  # (3, S, 1) - (3, 1, D) -> (3, S, D)
            dist = torch.linalg.norm(diff, dim=0)  # (S, D)
            closest_idx = torch.argmin(dist, dim=1)  # (S,)
            idx_list.append(closest_idx)
        closest_idx = torch.cat(idx_list)
        return closest_idx
//  [python]
def update_hess_q(
    hess_q: torch.Tensor,
    history: list["tico.opt.Step"],
    ic: tico.ic.IC,
    max_updates: int = 1,
) -> torch.Tensor:
    """Approximately update the Hessian matrix using the history of optimization steps.

    Args:
        hess_q: The current Hessian matrix.
        history: The history of optimization steps.
        ic: The internal coordinate system.
        max_updates: The maximum number of previous steps to use for the update.

    Returns:
        The updated Hessian matrix.
    """
    if len(history) < 2:
        return hess_q

    n_steps = 0

    for _ in range(2, len(history) + 1):
        if n_steps == max_updates:
            break
        n_steps += 1

    for i in range(n_steps):
        this_step = -n_steps + i
        prev_step = -n_steps + i - 1

        dq = ic.compute_dq(
            history[this_step].coords_x, history[prev_step].coords_x
        ).unsqueeze(-1)
        dg = (history[prev_step].grad_q - history[this_step].grad_q).unsqueeze(-1)

        if torch.norm(dq) < 1e-6:
            continue
        if torch.norm(dg) < 1e-6:
            continue

        mat_1 = (dg @ dg.T) / (dg.T @ dq)
        mat_2 = ((hess_q @ dq) @ (hess_q @ dq).T) / torch.linalg.multi_dot(
            [dq.T, hess_q, dq]
        )
        hess_q = hess_q + mat_1 - mat_2

    return hess_q
//  [python]
def update_hess_q(
    hess_q: torch.Tensor,
    history: list["tico.opt.Step"],
    ic: tico.ic.IC,
    max_updates: int = 1,
) -> torch.Tensor:
    """Approximately update the Hessian matrix using the history of optimization steps.

    Args:
        hess_q: The current Hessian matrix.
        history: The history of optimization steps.
        ic: The internal coordinate system.
        max_updates: The maximum number of previous steps to use for the update.

    Returns:
        The updated Hessian matrix.
    """
    if len(history) < 2:
        return hess_q

    n_steps = 0

    for _ in range(2, len(history) + 1):
        if n_steps == max_updates:
            break
        n_steps += 1

    for i in range(n_steps):
        this_step = -n_steps + i
        prev_step = -n_steps + i - 1

        dq = ic.compute_dq(
            history[this_step].coords_x, history[prev_step].coords_x
        ).unsqueeze(-1)
        dg = (history[prev_step].grad_q - history[this_step].grad_q).unsqueeze(-1)

        if torch.linalg.norm(dq) < 1e-6:
            continue
        if torch.linalg.norm(dg) < 1e-6:
            continue

        mat_1 = (dg @ dg.T) / (dg.T @ dq)
        mat_2 = ((hess_q @ dq) @ (hess_q @ dq).T) / torch.linalg.multi_dot(
            [dq.T, hess_q, dq]
        )
        hess_q = hess_q + mat_1 - mat_2

    return hess_q
//  [python]
def test_accuracy(tensor, ranks):
    """
    Given a tensor, this function returns the relative error derived from projecting such
    tensor using HOSVD and then going back to the reconstructed original tensor via
    the inverse projection.
    :param torch.Tensor tensor: the input tensor to be tested on
    :param list[int] ranks: list of the ranks of the individual directional projections
                            (len(ranks) must be equal to len(tensor.shape))
    :return: relative error of the method and list of the singular values of each unfolding
    :rtype: float, list[torch.Tensor]
    """
    HOSVD = hosvd(ranks)
    HOSVD.fit(tensor, return_S_tensor=True)
    red_tensor = HOSVD.transform(tensor)
    reconstruct_tensor = HOSVD.inverse_transform(red_tensor)
    error = torch.norm(reconstruct_tensor - tensor)
    relative_error = error / np.linalg.norm(tensor)
    return relative_error, HOSVD.modal_singular_values
//  [python]
def test_accuracy(tensor, ranks):
    """
    Given a tensor, this function returns the relative error derived from projecting such
    tensor using HOSVD and then going back to the reconstructed original tensor via
    the inverse projection.
    :param torch.Tensor tensor: the input tensor to be tested on
    :param list[int] ranks: list of the ranks of the individual directional projections
                            (len(ranks) must be equal to len(tensor.shape))
    :return: relative error of the method and list of the singular values of each unfolding
    :rtype: float, list[torch.Tensor]
    """
    HOSVD = hosvd(ranks)
    HOSVD.fit(tensor, return_S_tensor=True)
    red_tensor = HOSVD.transform(tensor)
    reconstruct_tensor = HOSVD.inverse_transform(red_tensor)
    error = torch.linalg.norm(reconstruct_tensor - tensor)
    relative_error = error / np.linalg.norm(tensor)
    return relative_error, HOSVD.modal_singular_values
//  [python]
def compute_pose_similarity(poses):
    # Pairwise scale align the poses before comparing them
    square_scales = torch.mean(torch.square(poses), dim=(-2, -1), keepdim=True)
    square_scales1 = torch.unsqueeze(square_scales, 0)
    square_scales2 = torch.unsqueeze(square_scales, 1)
    mean_square_scales = (square_scales1 + square_scales2) / 2
    scale_factor1 = torch.sqrt(mean_square_scales / square_scales1)
    scale_factor2 = torch.sqrt(mean_square_scales / square_scales2)

    poses1 = torch.unsqueeze(poses, 0)
    poses2 = torch.unsqueeze(poses, 1)

    dists = torch.norm(scale_factor1 * poses1 - scale_factor2 * poses2, dim=-1)
    best_dists = torch.topk(dists, k=poses.shape[-2] // 4, sorted=False).values
    return torch.mean(torch.relu(1 - best_dists / 300), dim=-1)
//  [python]
def compute_pose_similarity(poses):
    # Pairwise scale align the poses before comparing them
    square_scales = torch.mean(torch.square(poses), dim=(-2, -1), keepdim=True)
    square_scales1 = torch.unsqueeze(square_scales, 0)
    square_scales2 = torch.unsqueeze(square_scales, 1)
    mean_square_scales = (square_scales1 + square_scales2) / 2
    scale_factor1 = torch.sqrt(mean_square_scales / square_scales1)
    scale_factor2 = torch.sqrt(mean_square_scales / square_scales2)

    poses1 = torch.unsqueeze(poses, 0)
    poses2 = torch.unsqueeze(poses, 1)

    dists = torch.linalg.norm(scale_factor1 * poses1 - scale_factor2 * poses2, dim=-1)
    best_dists = torch.topk(dists, k=poses.shape[-2] // 4, sorted=False).values
    return torch.mean(torch.relu(1 - best_dists / 300), dim=-1)
//  [python]
def _normalize_camera(scene_cameras: PerspectiveCameras, distance=3.5, look_at_optical_center=False):
        '''
        Normalize the camera origins to be roughly distance 1.5 away from origin
        '''
        if look_at_optical_center:
            look_at_center = _get_nearest_centroid(scene_cameras).unsqueeze(-1).unsqueeze(0)
            scene_shift = (scene_cameras.R.permute(0,2,1)@(look_at_center))[:,:,0]
            scene_cameras = PerspectiveCameras(
                R = scene_cameras.R, 
                T = scene_cameras.T - scene_shift, 
                focal_length = scene_cameras.focal_length,
                principal_point = scene_cameras.principal_point,
                image_size = scene_cameras.image_size,
                device = scene_cameras.device,
            )
             
        cam_dist_mean = torch.mean(torch.norm(scene_cameras.get_camera_center(), dim=1))
        scale_factor = 1.0 / (cam_dist_mean / distance)

        normalized_cameras = PerspectiveCameras(
            R = scene_cameras.R, 
            T = scene_cameras.T * scale_factor, 
            focal_length = scene_cameras.focal_length,
            principal_point = scene_cameras.principal_point,
            image_size = scene_cameras.image_size,
            device = scene_cameras.device,
        )
        return normalized_cameras
//  [python]
def _normalize_camera(scene_cameras: PerspectiveCameras, distance=3.5, look_at_optical_center=False):
        '''
        Normalize the camera origins to be roughly distance 1.5 away from origin
        '''
        if look_at_optical_center:
            look_at_center = _get_nearest_centroid(scene_cameras).unsqueeze(-1).unsqueeze(0)
            scene_shift = (scene_cameras.R.permute(0,2,1)@(look_at_center))[:,:,0]
            scene_cameras = PerspectiveCameras(
                R = scene_cameras.R, 
                T = scene_cameras.T - scene_shift, 
                focal_length = scene_cameras.focal_length,
                principal_point = scene_cameras.principal_point,
                image_size = scene_cameras.image_size,
                device = scene_cameras.device,
            )
             
        cam_dist_mean = torch.mean(torch.linalg.norm(scene_cameras.get_camera_center(), axis=1))
        scale_factor = 1.0 / (cam_dist_mean / distance)

        normalized_cameras = PerspectiveCameras(
            R = scene_cameras.R, 
            T = scene_cameras.T * scale_factor, 
            focal_length = scene_cameras.focal_length,
            principal_point = scene_cameras.principal_point,
            image_size = scene_cameras.image_size,
            device = scene_cameras.device,
        )
        return normalized_cameras
//  [python]
def scale_emb_to_mag(base_emb, weighted_emb):
    norm_base = torch.norm(base_emb)
    norm_weighted = torch.norm(weighted_emb)
    embeddings_final = (norm_base / norm_weighted) * weighted_emb
    return embeddings_final
//  [python]
def scale_emb_to_mag(base_emb, weighted_emb):
    norm_base = torch.linalg.norm(base_emb)
    norm_weighted = torch.linalg.norm(weighted_emb)
    embeddings_final = (norm_base / norm_weighted) * weighted_emb
    return embeddings_final
//  [python]
def get_dist_to_ligand(pred_xyz, lig_xyz):
    '''
    Calculate distances between predicted BB coordinates (superimposed to template frame) 
    to all heavy atoms in ligand.

    Input:
        - pred_xyz: predicted coordinates in reference frame (B, L, 3, 3)
        - lig_xyz:  ligand coordinates in reference frame (B, L_lig*3, 3)
    Output:
        - dist_lig: distances from predicted motif to ligand
    '''
    pred_xyz = pred_xyz[:,:,:3,:]
    B, L     = pred_xyz.shape[:2]
    pred_xyz = pred_xyz.reshape(B, L*3, 3)

    lig_xyz = lig_xyz[None].to(pred_xyz.device)
    #L_lig = lig_xyz.shape[1]
    #lig_xyz = lig_xyz.view(1, L_lig*3, 3)

    # calculate distances
    d = lig_xyz[:,:,None,:] - pred_xyz[:,None,:,:]
    dist = torch.norm(d,dim=3)
    return dist
//  [python]
def get_dist_to_ligand(pred_xyz, lig_xyz):
    '''
    Calculate distances between predicted BB coordinates (superimposed to template frame) 
    to all heavy atoms in ligand.

    Input:
        - pred_xyz: predicted coordinates in reference frame (B, L, 3, 3)
        - lig_xyz:  ligand coordinates in reference frame (B, L_lig*3, 3)
    Output:
        - dist_lig: distances from predicted motif to ligand
    '''
    pred_xyz = pred_xyz[:,:,:3,:]
    B, L     = pred_xyz.shape[:2]
    pred_xyz = pred_xyz.reshape(B, L*3, 3)

    lig_xyz = lig_xyz[None].to(pred_xyz.device)
    #L_lig = lig_xyz.shape[1]
    #lig_xyz = lig_xyz.view(1, L_lig*3, 3)

    # calculate distances
    d = lig_xyz[:,:,None,:] - pred_xyz[:,None,:,:]
    dist = torch.linalg.norm(d,dim=3)
    return dist
//  [python]
def forward(self, x):
        """"""
        x = torch.nn.functional.softmax(x, dim=1)
        return -torch.norm(x, "nuc") / x.shape[0]
//  [python]
def forward(self, x):
        """"""
        x = torch.nn.functional.softmax(x, dim=1)
        return -torch.linalg.norm(x, "nuc") / x.shape[0]
//  [python]
def compare_traj(ref_traj, states_reached, robot_model, experiment_name=""):
    """Compute and print error between state log and the reference trajectory"""
    # Compute error
    pos_err = torch.zeros(len(states_reached))
    ori_err = torch.zeros(len(states_reached))
    for i in range(len(states_reached)):
        joint_pos_actual = torch.Tensor(states_reached[i][1].joint_positions)
        pose_actual = torch.cat(robot_model.forward_kinematics(joint_pos_actual))
        pose_desired = ref_traj[states_reached[i][0]]

        pos_diff = pose_desired[:3] - pose_actual[:3]
        ori_diff = R.from_quat(pose_actual[3:]).inv() * R.from_quat(pose_desired[3:])
        pos_err[i] = torch.norm(pos_diff)
        ori_err[i] = torch.norm(ori_diff.as_rotvec())

    # Convert, compute, and print stats
    pos_err_mm = pos_err * 1000.0
    ori_err_deg = ori_err * 180 / np.pi

    pos_err_mean = torch.mean(pos_err_mm)
    pos_err_std = torch.std(pos_err_mm)
    pos_err_max = torch.max(pos_err_mm)
    ori_err_mean = torch.mean(ori_err_deg)
    ori_err_std = torch.std(ori_err_deg)
    ori_err_max = torch.max(ori_err_deg)

    print(f"=== {experiment_name} tracking results ===")
    print(
        f"\tPos error(mm): mean={pos_err_mean:.2f}, std={pos_err_std:.2f}, max={pos_err_max:.2f}"
    )
    print(
        f"\tOri error(deg): mean={ori_err_mean:.2f}, std={ori_err_std:.2f}, max={ori_err_max:.2f}"
    )
//  [python]
def compare_traj(ref_traj, states_reached, robot_model, experiment_name=""):
    """Compute and print error between state log and the reference trajectory"""
    # Compute error
    pos_err = torch.zeros(len(states_reached))
    ori_err = torch.zeros(len(states_reached))
    for i in range(len(states_reached)):
        joint_pos_actual = torch.Tensor(states_reached[i][1].joint_positions)
        pose_actual = torch.cat(robot_model.forward_kinematics(joint_pos_actual))
        pose_desired = ref_traj[states_reached[i][0]]

        pos_diff = pose_desired[:3] - pose_actual[:3]
        ori_diff = R.from_quat(pose_actual[3:]).inv() * R.from_quat(pose_desired[3:])
        pos_err[i] = torch.linalg.norm(pos_diff)
        ori_err[i] = torch.linalg.norm(ori_diff.as_rotvec())

    # Convert, compute, and print stats
    pos_err_mm = pos_err * 1000.0
    ori_err_deg = ori_err * 180 / np.pi

    pos_err_mean = torch.mean(pos_err_mm)
    pos_err_std = torch.std(pos_err_mm)
    pos_err_max = torch.max(pos_err_mm)
    ori_err_mean = torch.mean(ori_err_deg)
    ori_err_std = torch.std(ori_err_deg)
    ori_err_max = torch.max(ori_err_deg)

    print(f"=== {experiment_name} tracking results ===")
    print(
        f"\tPos error(mm): mean={pos_err_mean:.2f}, std={pos_err_std:.2f}, max={pos_err_max:.2f}"
    )
    print(
        f"\tOri error(deg): mean={ori_err_mean:.2f}, std={ori_err_std:.2f}, max={ori_err_max:.2f}"
    )
//  [python]
def lagrangian(self, batch_int,eval_args, items_to_sample, alpha, slack, lam, mu,
                   eta=settings.ETA_LAGRANGIAN_DEFAULT, kappa=settings.KAPPA_LAGRANGIAN_DEFAULT):
        F = self.train_objective(
            batch_int,eval_args=eval_args, items_to_sample=items_to_sample)
        H = self.train_constraint(batch_int,eval_args=eval_args, items_to_sample=items_to_sample,
                                  alpha=alpha, slack=slack, eta=eta, kappa=kappa)
        return F + lam @ H + (mu/2)*(t.norm(H)**2), H.detach()
//  [python]
def lagrangian(self, batch_int,eval_args, items_to_sample, alpha, slack, lam, mu,
                   eta=settings.ETA_LAGRANGIAN_DEFAULT, kappa=settings.KAPPA_LAGRANGIAN_DEFAULT):
        F = self.train_objective(
            batch_int,eval_args=eval_args, items_to_sample=items_to_sample)
        H = self.train_constraint(batch_int,eval_args=eval_args, items_to_sample=items_to_sample,
                                  alpha=alpha, slack=slack, eta=eta, kappa=kappa)
        return F + lam @ H + (mu/2)*(torch.linalg.norm(H)**2), H.detach()
//  [python]
def _get_relative_camera(scene_cameras:PerspectiveCameras, query_idx=None, query_camera=None, center_at_origin=False, shift_z=False):
        """
        Transform context cameras relative to a base query camera
        Args:
            scene_cameras (PyTorch3D Camera): cameras object
            query_idx (tensor or List): a length 1 list defining query idx
        Returns:
            cams_relative (PyTorch3D Camera): cameras object relative to query camera
        """
        if center_at_origin is False:
            assert shift_z is False

        if query_idx is not None:
            query_camera = _get_camera_slice(scene_cameras, query_idx)
        else:
            assert query_camera is not None

        if shift_z:
            optical_center = _get_nearest_centroid(scene_cameras).to(scene_cameras.device)
            z_dist = torch.norm(query_camera.get_camera_center()[0] - optical_center.unsqueeze(0), dim=1)

        
        query_world2view = query_camera.get_world_to_view_transform()
        all_world2view = scene_cameras.get_world_to_view_transform()
        
        if center_at_origin:
            identity_cam = PerspectiveCameras(device=scene_cameras.device, R=query_camera.R, T=query_camera.T)
        else:
            T = torch.zeros((1, 3))
            identity_cam = PerspectiveCameras(device=scene_cameras.device, R=query_camera.R, T=T)
         
        identity_world2view  = identity_cam.get_world_to_view_transform()

        # compose the relative transformation as g_i^{-1} g_j
        relative_world2view = identity_world2view.inverse().compose(all_world2view)
        
        # generate a camera from the relative transform
        relative_matrix = relative_world2view.get_matrix()
        cams_relative = PerspectiveCameras(
                            R = relative_matrix[:, :3, :3],
                            T = relative_matrix[:, 3, :3],
                            focal_length = scene_cameras.focal_length,
                            principal_point = scene_cameras.principal_point,
                            image_size = scene_cameras.image_size,
                            device = scene_cameras.device,
                        )
        
        if shift_z:
            scene_shift = (cams_relative.R.permute(0,2,1)@(torch.tensor([[0.0],[0.0],[z_dist]], device=scene_cameras.device)))[:,:,0]
            cams_relative = PerspectiveCameras(
                R = cams_relative.R, 
                T = cams_relative.T + scene_shift, 
                focal_length = cams_relative.focal_length,
                principal_point = cams_relative.principal_point,
                image_size = cams_relative.image_size,
                device = cams_relative.device,
            )
        return cams_relative
//  [python]
def _get_relative_camera(scene_cameras:PerspectiveCameras, query_idx=None, query_camera=None, center_at_origin=False, shift_z=False):
        """
        Transform context cameras relative to a base query camera
        Args:
            scene_cameras (PyTorch3D Camera): cameras object
            query_idx (tensor or List): a length 1 list defining query idx
        Returns:
            cams_relative (PyTorch3D Camera): cameras object relative to query camera
        """
        if center_at_origin is False:
            assert shift_z is False

        if query_idx is not None:
            query_camera = _get_camera_slice(scene_cameras, query_idx)
        else:
            assert query_camera is not None

        if shift_z:
            optical_center = _get_nearest_centroid(scene_cameras).to(scene_cameras.device)
            z_dist = torch.linalg.norm(query_camera.get_camera_center()[0] - optical_center.unsqueeze(0), axis=1)

        
        query_world2view = query_camera.get_world_to_view_transform()
        all_world2view = scene_cameras.get_world_to_view_transform()
        
        if center_at_origin:
            identity_cam = PerspectiveCameras(device=scene_cameras.device, R=query_camera.R, T=query_camera.T)
        else:
            T = torch.zeros((1, 3))
            identity_cam = PerspectiveCameras(device=scene_cameras.device, R=query_camera.R, T=T)
         
        identity_world2view  = identity_cam.get_world_to_view_transform()

        # compose the relative transformation as g_i^{-1} g_j
        relative_world2view = identity_world2view.inverse().compose(all_world2view)
        
        # generate a camera from the relative transform
        relative_matrix = relative_world2view.get_matrix()
        cams_relative = PerspectiveCameras(
                            R = relative_matrix[:, :3, :3],
                            T = relative_matrix[:, 3, :3],
                            focal_length = scene_cameras.focal_length,
                            principal_point = scene_cameras.principal_point,
                            image_size = scene_cameras.image_size,
                            device = scene_cameras.device,
                        )
        
        if shift_z:
            scene_shift = (cams_relative.R.permute(0,2,1)@(torch.tensor([[0.0],[0.0],[z_dist]], device=scene_cameras.device)))[:,:,0]
            cams_relative = PerspectiveCameras(
                R = cams_relative.R, 
                T = cams_relative.T + scene_shift, 
                focal_length = cams_relative.focal_length,
                principal_point = cams_relative.principal_point,
                image_size = cams_relative.image_size,
                device = cams_relative.device,
            )
        return cams_relative
//  [python]
def normalize_quat(quats):
        dtype = quats.dtype
        quats = quats.float()
        quats = quats / torch.norm(quats, dim=-1, keepdim=True)
        quats = quats.type(dtype)
        return quats
//  [python]
def normalize_quat(quats):
        dtype = quats.dtype
        quats = quats.float()
        quats = quats / torch.linalg.norm(quats, dim=-1, keepdim=True)
        quats = quats.type(dtype)
        return quats
//  [python]
def test_contrastive_loss():
    """ Test tensor contrastive loss """

    def contrast(df: float, dy: float, threshold: float) -> float:
        """ Contrastive loss for continuous labels """
        if dy > threshold:
            s = (2 - min(threshold, df) / threshold) * (dy - max(df, threshold))
        else:
            s = max(threshold, df) / threshold * (min(df, threshold) - dy)
        return max(0, s)

    ys = torch.rand(10, 1)
    embs = torch.rand(10, 3)
    constr_loss = ContrastiveLossTorch(.25)

    loss = constr_loss.build_loss_matrix(embs, ys)

    for i in range(loss.shape[0]):
        for j in range(loss.shape[1]):
            df_ = torch.norm(embs[i] - embs[j]).item()
            dy_ = torch.abs(ys[i] - ys[j]).item()
            if not np.isclose(loss[i, j].item(), contrast(df_, dy_, constr_loss.threshold)):
                return 0

    print(f"Passed: test_contrastive_loss")
    return 1
//  [python]
def test_contrastive_loss():
    """ Test tensor contrastive loss """

    def contrast(df: float, dy: float, threshold: float) -> float:
        """ Contrastive loss for continuous labels """
        if dy > threshold:
            s = (2 - min(threshold, df) / threshold) * (dy - max(df, threshold))
        else:
            s = max(threshold, df) / threshold * (min(df, threshold) - dy)
        return max(0, s)

    ys = torch.rand(10, 1)
    embs = torch.rand(10, 3)
    constr_loss = ContrastiveLossTorch(.25)

    loss = constr_loss.build_loss_matrix(embs, ys)

    for i in range(loss.shape[0]):
        for j in range(loss.shape[1]):
            df_ = torch.linalg.norm(embs[i] - embs[j]).item()
            dy_ = torch.abs(ys[i] - ys[j]).item()
            if not np.isclose(loss[i, j].item(), contrast(df_, dy_, constr_loss.threshold)):
                return 0

    print(f"Passed: test_contrastive_loss")
    return 1
//  [python]
def norm_ft(X: torch.Tensor, ord: Optional[Union[int, float]] = None) -> torch.Tensor:
    r"""Normalize the input feature matrix with specified ``ord`` refer to pytorch's `torch.linalg.norm <https://pytorch.org/docs/stable/generated/torch.linalg.norm.html#torch.linalg.norm>`_ function.

    .. note::
        The input feature matrix is expected to be a 1D vector or a 2D tensor with shape (num_samples, num_features).

    Args:
        ``X`` (``torch.Tensor``): The input feature.
        ``ord`` (``Union[int, float]``, optional): The order of the norm can be either an ``int``, ``float``. If ``ord`` is ``None``, the norm is computed with the 2-norm. Defaults to ``None``.

    Examples:
        >>> import easygraph.datapipe as dd
        >>> import torch
        >>> X = torch.tensor([
                    [0.1, 0.2, 0.5],
                    [0.5, 0.2, 0.3],
                    [0.3, 0.2, 0]
                ])
        >>> dd.norm_ft(X)
        tensor([[0.1826, 0.3651, 0.9129],
                [0.8111, 0.3244, 0.4867],
                [0.8321, 0.5547, 0.0000]])
    """
    if X.dim() == 1:
        X_norm = 1 / torch.norm(X, p=ord)
        X_norm[torch.isinf(X_norm)] = 0
        return X * X_norm
    elif X.dim() == 2:
        X_norm = 1 / torch.norm(X, p=ord, dim=1, keepdim=True)
        X_norm[torch.isinf(X_norm)] = 0
        return X * X_norm
    else:
        raise ValueError(
            "The input feature matrix is expected to be a 1D verter or a 2D tensor with"
            " shape (num_samples, num_features)."
        )
//  [python]
def norm_ft(X: torch.Tensor, ord: Optional[Union[int, float]] = None) -> torch.Tensor:
    r"""Normalize the input feature matrix with specified ``ord`` refer to pytorch's `torch.linalg.norm <https://pytorch.org/docs/stable/generated/torch.linalg.norm.html#torch.linalg.norm>`_ function.

    .. note::
        The input feature matrix is expected to be a 1D vector or a 2D tensor with shape (num_samples, num_features).

    Args:
        ``X`` (``torch.Tensor``): The input feature.
        ``ord`` (``Union[int, float]``, optional): The order of the norm can be either an ``int``, ``float``. If ``ord`` is ``None``, the norm is computed with the 2-norm. Defaults to ``None``.

    Examples:
        >>> import easygraph.datapipe as dd
        >>> import torch
        >>> X = torch.tensor([
                    [0.1, 0.2, 0.5],
                    [0.5, 0.2, 0.3],
                    [0.3, 0.2, 0]
                ])
        >>> dd.norm_ft(X)
        tensor([[0.1826, 0.3651, 0.9129],
                [0.8111, 0.3244, 0.4867],
                [0.8321, 0.5547, 0.0000]])
    """
    if X.dim() == 1:
        X_norm = 1 / torch.linalg.norm(X, ord=ord)
        X_norm[torch.isinf(X_norm)] = 0
        return X * X_norm
    elif X.dim() == 2:
        X_norm = 1 / torch.linalg.norm(X, ord=ord, dim=1, keepdim=True)
        X_norm[torch.isinf(X_norm)] = 0
        return X * X_norm
    else:
        raise ValueError(
            "The input feature matrix is expected to be a 1D verter or a 2D tensor with"
            " shape (num_samples, num_features)."
        )
//  [python]
def test_circle_random_sampling_with_n():
    C = Circle(R2('x'), [0, 0], 4)
    points = C.sample_random_uniform(n=40)
    assert points.as_tensor.shape == (40, 2)
    assert all(torch.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_circle_random_sampling_with_n():
    C = Circle(R2('x'), [0, 0], 4)
    points = C.sample_random_uniform(n=40)
    assert points.as_tensor.shape == (40, 2)
    assert all(torch.linalg.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_circle_random_sampling_with_d():
    C = Circle(R2('x'), [0, 0], 4)
    points = C.sample_random_uniform(d=13)
    assert all(torch.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_circle_random_sampling_with_d():
    C = Circle(R2('x'), [0, 0], 4)
    points = C.sample_random_uniform(d=13)
    assert all(torch.linalg.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def update_from_gradients(self):
        self.aggregated_gradients = [
            torch.zeros_like(params) for params in self.server_model.parameters()
        ]
        grad_norms = []

        M = len(self.uploaded_gradients)
        len_gradients = len(self.aggregated_gradients)

        for gradients in self.uploaded_gradients:
            for gradient_id in range(len_gradients):
                self.aggregated_gradients[gradient_id] = (
                    gradients[gradient_id] * (1 / M)
                    + self.aggregated_gradients[gradient_id]
                )
                grad_norms.append(
                    sum([g.norm() ** 2 for g in gradients[gradient_id]])
                )

        agg_grad_norm = sum(
            [g.norm() ** 2 for g in self.aggregated_gradients]
        )

        self.optimizer.lr = max(
            1, sum([g / (2 * M * (agg_grad_norm + self.eps)) for g in grad_norms])
        )
        self.optimizer.step(self.aggregated_gradients)
//  [python]
def update_from_gradients(self):
        self.aggregated_gradients = [
            torch.zeros_like(params) for params in self.server_model.parameters()
        ]
        grad_norms = []

        M = len(self.uploaded_gradients)
        len_gradients = len(self.aggregated_gradients)

        for gradients in self.uploaded_gradients:
            for gradient_id in range(len_gradients):
                self.aggregated_gradients[gradient_id] = (
                    gradients[gradient_id] * (1 / M)
                    + self.aggregated_gradients[gradient_id]
                )
                grad_norms.append(
                    sum([torch.linalg.norm(g) ** 2 for g in gradients[gradient_id]])
                )

        agg_grad_norm = sum(
            [torch.linalg.norm(g) ** 2 for g in self.aggregated_gradients]
        )

        self.optimizer.lr = max(
            1, sum([g / (2 * M * (agg_grad_norm + self.eps)) for g in grad_norms])
        )
        self.optimizer.step(self.aggregated_gradients)
//  [python]
def test_circle_grid_sampling_with_n():
    C = Circle(R2('x'), [0, 0], 4)
    points = C.sample_grid(n=40)
    assert points.as_tensor.shape == (40, 2)
    assert all(torch.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def test_circle_grid_sampling_with_n():
    C = Circle(R2('x'), [0, 0], 4)
    points = C.sample_grid(n=40)
    assert points.as_tensor.shape == (40, 2)
    assert all(torch.linalg.norm(points.as_tensor, dim=1) <= 4.0)
//  [python]
def __init__(
        self,
        rot_mats: Optional[torch.Tensor] = None,
        quats: Optional[torch.Tensor] = None,
        normalize_quats: bool = True,
    ):
        if (rot_mats is None and quats is None) or (
            rot_mats is not None and quats is not None
        ):
            raise ValueError("Exactly one input argument must be specified")

        if (rot_mats is not None and rot_mats.shape[-2:] != (3, 3)) or (
            quats is not None and quats.shape[-1] != 4
        ):
            raise ValueError(
                "Incorrectly shaped rotation matrix or quaternion"
            )

        # Force full-precision
        if quats is not None:
            quats = quats.to(dtype=torch.float32)
        if rot_mats is not None:
            rot_mats = rot_mats.to(dtype=torch.float32)

        if quats is not None and normalize_quats:
            quats = quats / torch.norm(quats, dim=-1, keepdim=True)

        self._rot_mats = rot_mats
        self._quats = quats