jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_norm_dtype(self, device, dtype):
        make_arg = partial(make_tensor, dtype=dtype, device=device)

        def run_test_case(input_size, ord, keepdim, to_dtype):
            msg = (
                f'input_size={input_size}, ord={ord}, keepdim={keepdim}, '
                f'dtype={dtype}, to_dtype={to_dtype}')
            input = make_arg(input_size)
            result = torch.linalg.norm(input, ord, keepdim=keepdim)
            self.assertEqual(result.dtype, input.real.dtype, msg=msg)

            result_out = torch.empty((0), dtype=result.dtype, device=device)
            torch.linalg.norm(input, ord, keepdim=keepdim, out=result_out)
            self.assertEqual(result, result_out, msg=msg)

            result = torch.linalg.norm(input.to(to_dtype), ord, keepdim=keepdim)
            result_with_dtype = torch.linalg.norm(input, ord, keepdim=keepdim, dtype=to_dtype)
            self.assertEqual(result, result_with_dtype, msg=msg)

            result_out_with_dtype = torch.empty_like(result_with_dtype)
            torch.linalg.norm(input, ord, keepdim=keepdim, dtype=to_dtype, out=result_out_with_dtype)
            self.assertEqual(result_with_dtype, result_out_with_dtype, msg=msg)

        ord_vector = [0, 1, -1, 2, -2, 3, -3, 4.5, -4.5, inf, -inf, None]

        # In these orders we are computing the 10-th power and 10-th root of numbers.
        # We avoid them for half-precision types as it makes the tests above too badly conditioned
        if dtype != torch.float16 and dtype != torch.bfloat16:
            ord_vector.extend([0.1, -0.1])
        ord_matrix = ['fro', 'nuc', 1, -1, 2, -2, inf, -inf, None]
        S = 10

        if dtype == torch.cfloat:
            norm_dtypes = (torch.cfloat, torch.cdouble)
        elif dtype == torch.cdouble:
            norm_dtypes = (torch.cdouble,)
        elif dtype in (torch.float16, torch.bfloat16, torch.float):
            norm_dtypes = (torch.float, torch.double)
        elif dtype == torch.double:
            norm_dtypes = (torch.double,)
        else:
            raise RuntimeError("Unsupported dtype")

        for ord, keepdim, norm_dtype in product(ord_vector, (True, False), norm_dtypes):
            run_test_case((S,) , ord, keepdim, norm_dtype)

        for ord, keepdim, norm_dtype in product(ord_matrix, (True, False), norm_dtypes):
            if ord in [2, -2, 'nuc']:
                # We need torch.svdvals
                if dtype == torch.float16 or dtype == torch.bfloat16:
                    continue

                # We need LAPACK or equivalent
                if ((torch.device(device).type == 'cuda' and not torch.cuda.has_magma and not has_cusolver()) or
                   (torch.device(device).type == 'cpu' and not torch._C.has_lapack)):
                    continue
            run_test_case((S, S) , ord, keepdim, norm_dtype)
//  [python]
def sample_hypersphere(
    d: int,
    n: int = 1,
    qmc: bool = False,
    seed: Optional[int] = None,
    device: Optional[torch.device] = None,
    dtype: Optional[torch.dtype] = None,
) -> Tensor:
    r"""Sample uniformly from a unit d-sphere.

    Args:
        d: The dimension of the hypersphere.
        n: The number of samples to return.
        qmc: If True, use QMC Sobol sampling (instead of i.i.d. uniform).
        seed: If provided, use as a seed for the RNG.
        device: The torch device.
        dtype:  The torch dtype.

    Returns:
        An  `n x d` tensor of uniform samples from from the d-hypersphere.

    Example:
        >>> sample_hypersphere(d=5, n=10)
    """
    dtype = torch.float if dtype is None else dtype
    if d == 1:
        rnd = torch.randint(0, 2, (n, 1), device=device, dtype=dtype)
        return 2 * rnd - 1
    if qmc:
        rnd = draw_sobol_normal_samples(d=d, n=n, device=device, dtype=dtype, seed=seed)
    else:
        with manual_seed(seed=seed):
            rnd = torch.randn(n, d, dtype=dtype)
    samples = rnd / t.norm(rnd, dim=-1, keepdim=True)
    if device is not None:
        samples = samples.to(device)
    return samples
//  [python]
def sample_hypersphere(
    d: int,
    n: int = 1,
    qmc: bool = False,
    seed: Optional[int] = None,
    device: Optional[torch.device] = None,
    dtype: Optional[torch.dtype] = None,
) -> Tensor:
    r"""Sample uniformly from a unit d-sphere.

    Args:
        d: The dimension of the hypersphere.
        n: The number of samples to return.
        qmc: If True, use QMC Sobol sampling (instead of i.i.d. uniform).
        seed: If provided, use as a seed for the RNG.
        device: The torch device.
        dtype:  The torch dtype.

    Returns:
        An  `n x d` tensor of uniform samples from from the d-hypersphere.

    Example:
        >>> sample_hypersphere(d=5, n=10)
    """
    dtype = torch.float if dtype is None else dtype
    if d == 1:
        rnd = torch.randint(0, 2, (n, 1), device=device, dtype=dtype)
        return 2 * rnd - 1
    if qmc:
        rnd = draw_sobol_normal_samples(d=d, n=n, device=device, dtype=dtype, seed=seed)
    else:
        with manual_seed(seed=seed):
            rnd = torch.randn(n, d, dtype=dtype)
    samples = rnd / torch.linalg.norm(rnd, dim=-1, keepdim=True)
    if device is not None:
        samples = samples.to(device)
    return samples
//  [python]
def Broyden_step(J, h, Yp, Yph):
    delta = torch.matmul(J, h)
    # avoid constructing a second giant jacobian matrix, instead go one row at a time
    for j in range(J.shape[1]):
        J[:, j] += (Yph - Yp - delta) * h[j] / torch.norm(h)
    return J
//  [python]
def Broyden_step(J, h, Yp, Yph):
    delta = torch.matmul(J, h)
    # avoid constructing a second giant jacobian matrix, instead go one row at a time
    for j in range(J.shape[1]):
        J[:, j] += (Yph - Yp - delta) * h[j] / torch.linalg.norm(h)
    return J
//  [python]
def pers2w(self, point_xyz_pers, camrotc2w, campos):
        #     point_xyz_pers    B X M X 3

        x_pers = point_xyz_pers[..., 0] * point_xyz_pers[..., 2]
        y_pers = point_xyz_pers[..., 1] * point_xyz_pers[..., 2]
        z_pers = point_xyz_pers[..., 2]
        xyz_c = torch.stack([x_pers, y_pers, z_pers], dim=-1)
        xyz_w_shift = torch.sum(xyz_c[...,None,:] * camrotc2w, dim=-1)
        # print("point_xyz_pers[..., 0, 0]", point_xyz_pers[..., 0, 0].shape, point_xyz_pers[..., 0, 0])
        ray_dirs = xyz_w_shift / (torch.norm(xyz_w_shift, dim=-1, keepdims=True) + 1e-7)

        xyz_w = xyz_w_shift + campos[:, None, :]
        return xyz_w, ray_dirs
//  [python]
def pers2w(self, point_xyz_pers, camrotc2w, campos):
        #     point_xyz_pers    B X M X 3

        x_pers = point_xyz_pers[..., 0] * point_xyz_pers[..., 2]
        y_pers = point_xyz_pers[..., 1] * point_xyz_pers[..., 2]
        z_pers = point_xyz_pers[..., 2]
        xyz_c = torch.stack([x_pers, y_pers, z_pers], dim=-1)
        xyz_w_shift = torch.sum(xyz_c[...,None,:] * camrotc2w, dim=-1)
        # print("point_xyz_pers[..., 0, 0]", point_xyz_pers[..., 0, 0].shape, point_xyz_pers[..., 0, 0])
        ray_dirs = xyz_w_shift / (torch.linalg.norm(xyz_w_shift, dim=-1, keepdims=True) + 1e-7)

        xyz_w = xyz_w_shift + campos[:, None, :]
        return xyz_w, ray_dirs
//  [python]
def barycentric_coordinate(pts, face_vertices):
    """
    :param pts: (B, N, 3)
    :param face_vertices: (B, N, 3, 3)
    :return bc_coords: (B, N, 3)
    """
    vec0 = face_vertices[:, :, 0] - pts
    vec1 = face_vertices[:, :, 1] - pts
    vec2 = face_vertices[:, :, 2] - pts
    area0 = torch.norm(torch.cross(vec1, vec2), dim = -1)
    area1 = torch.norm(torch.cross(vec2, vec0), dim = -1)
    area2 = torch.norm(torch.cross(vec0, vec1), dim = -1)
    bc_coord = torch.stack([area0, area1, area2], -1)
    bc_coord = F.normalize(bc_coord, p = 1, dim = -1, eps = 1e-16)
    return bc_coord
//  [python]
def barycentric_coordinate(pts, face_vertices):
    """
    :param pts: (B, N, 3)
    :param face_vertices: (B, N, 3, 3)
    :return bc_coords: (B, N, 3)
    """
    vec0 = face_vertices[:, :, 0] - pts
    vec1 = face_vertices[:, :, 1] - pts
    vec2 = face_vertices[:, :, 2] - pts
    area0 = torch.linalg.norm(torch.cross(vec1, vec2), dim = -1)
    area1 = torch.linalg.norm(torch.cross(vec2, vec0), dim = -1)
    area2 = torch.linalg.norm(torch.cross(vec0, vec1), dim = -1)
    bc_coord = torch.stack([area0, area1, area2], -1)
    bc_coord = F.normalize(bc_coord, p = 1, dim = -1, eps = 1e-16)
    return bc_coord
//  [python]
def sh_intrp(self, embedding, dists, pnt_mask, vsize, grid_vox_sz, axis_weight=None):
        # dists: B * R * SR * K * channel
        dist_norm = torch.norm(dists, dim=-1)
        dist_dirs = dists / torch.clamp(dist_norm[...,None], min=1e-8)
        shall = self.shcomp.sh_all(dist_dirs, filp_dir=False).view(dists.shape[:-1]+(self.shcomp.total_deg ** 2,))
        sh_coefs = embedding[..., :self.shcomp.total_deg ** 2]
        # shall: [1, 816, 24, 32, 16], sh_coefs: [1, 816, 24, 32, 16], pnt_mask: [1, 816, 24, 32]
        # debug: weights = pnt_mask * torch.sum(shall, dim=-1)
        # weights = pnt_mask * torch.sum(shall * getattr(self, self.opt.sh_act, None)(sh_coefs), dim=-1) * getattr(self, self.opt.sh_dist_func, None)(dist_norm)
        weights = pnt_mask * torch.sum(getattr(self, self.opt.sh_act, None)(shall * sh_coefs), dim=-1) * getattr(self, self.opt.sh_dist_func, None)(dist_norm) # changed
        return weights, embedding[..., self.shcomp.total_deg ** 2:]
//  [python]
def sh_intrp(self, embedding, dists, pnt_mask, vsize, grid_vox_sz, axis_weight=None):
        # dists: B * R * SR * K * channel
        dist_norm = torch.linalg.norm(dists, dim=-1)
        dist_dirs = dists / torch.clamp(dist_norm[...,None], min=1e-8)
        shall = self.shcomp.sh_all(dist_dirs, filp_dir=False).view(dists.shape[:-1]+(self.shcomp.total_deg ** 2,))
        sh_coefs = embedding[..., :self.shcomp.total_deg ** 2]
        # shall: [1, 816, 24, 32, 16], sh_coefs: [1, 816, 24, 32, 16], pnt_mask: [1, 816, 24, 32]
        # debug: weights = pnt_mask * torch.sum(shall, dim=-1)
        # weights = pnt_mask * torch.sum(shall * getattr(self, self.opt.sh_act, None)(sh_coefs), dim=-1) * getattr(self, self.opt.sh_dist_func, None)(dist_norm)
        weights = pnt_mask * torch.sum(getattr(self, self.opt.sh_act, None)(shall * sh_coefs), dim=-1) * getattr(self, self.opt.sh_dist_func, None)(dist_norm) # changed
        return weights, embedding[..., self.shcomp.total_deg ** 2:]
//  [python]
def point_robot_distance(link_spheres_tensor, points):
    """Compute distance between robot and points

    Args:
        link_spheres_tensor: [batch_robot, n_robot_spheres, 4]
        points: [batch_points, n_points, 3]

    Returns:
        distance: [batch_points, n_points]
    """
    if link_spheres_tensor.shape[0] != 1:
        assert link_spheres_tensor.shape[0] == points.shape[0]
    squeeze_shape = False
    n = 1
    if len(points.shape) == 2:
        squeeze_shape = True
        n, _ = points.shape
        points = points.unsqueeze(0)

    robot_spheres = link_spheres_tensor.view(link_spheres_tensor.shape[0], -1, 4).contiguous()
    robot_spheres = robot_spheres.unsqueeze(-3)

    robot_radius = robot_spheres[..., 3]
    points = points.unsqueeze(-2)
    sph_distance = -1 * (
        torch.norm(points - robot_spheres[..., :3], dim=-1) - robot_radius
    )  # b, n_spheres
    pt_distance = torch.max(sph_distance, dim=-1)[0]

    if squeeze_shape:
        pt_distance = pt_distance.view(n)
    return pt_distance
//  [python]
def point_robot_distance(link_spheres_tensor, points):
    """Compute distance between robot and points

    Args:
        link_spheres_tensor: [batch_robot, n_robot_spheres, 4]
        points: [batch_points, n_points, 3]

    Returns:
        distance: [batch_points, n_points]
    """
    if link_spheres_tensor.shape[0] != 1:
        assert link_spheres_tensor.shape[0] == points.shape[0]
    squeeze_shape = False
    n = 1
    if len(points.shape) == 2:
        squeeze_shape = True
        n, _ = points.shape
        points = points.unsqueeze(0)

    robot_spheres = link_spheres_tensor.view(link_spheres_tensor.shape[0], -1, 4).contiguous()
    robot_spheres = robot_spheres.unsqueeze(-3)

    robot_radius = robot_spheres[..., 3]
    points = points.unsqueeze(-2)
    sph_distance = -1 * (
        torch.linalg.norm(points - robot_spheres[..., :3], dim=-1) - robot_radius
    )  # b, n_spheres
    pt_distance = torch.max(sph_distance, dim=-1)[0]

    if squeeze_shape:
        pt_distance = pt_distance.view(n)
    return pt_distance
//  [python]
def nonlinear(self, Cz):
        norm = torch.norm(Cz.reshape(Cz.shape[0], Cz.shape[1], -1), dim=2)
        norm = torch.clamp(norm, min=1e-8)
        pred = tF.softmax(-self.lmbda * norm, dim=0).unsqueeze(2)
        y = torch.argmax(pred, axis=0) #TODO: for non argmax case
        gam = self.gam.unsqueeze(-1).unsqueeze(-1)
        out = torch.sum(gam * Cz * pred, axis=0)
        return out, y
//  [python]
def nonlinear(self, Cz):
        norm = torch.linalg.norm(Cz.reshape(Cz.shape[0], Cz.shape[1], -1), axis=2)
        norm = torch.clamp(norm, min=1e-8)
        pred = tF.softmax(-self.lmbda * norm, dim=0).unsqueeze(2)
        y = torch.argmax(pred, axis=0) #TODO: for non argmax case
        gam = self.gam.unsqueeze(-1).unsqueeze(-1)
        out = torch.sum(gam * Cz * pred, axis=0)
        return out, y
//  [python]
def normal_func(xyz: torch.Tensor, normal: torch.Tensor, sensor: torch.Tensor):
    assert normal is None, "normal already exists"
    assert sensor is not None, "please provide sensor positions for consistent orientations"

    xyz_numpy = xyz.cpu().numpy()
    indices, normal = pcu.estimate_point_cloud_normals_knn(xyz_numpy, 64)
    normal = torch.from_numpy(normal).to(xyz)
    indices = torch.from_numpy(indices).to(xyz).long()

    xyz, sensor = xyz[indices], sensor[indices]

    view_dir = sensor - xyz
    view_dir = view_dir / (torch.norm(view_dir, dim=-1, keepdim=True) + 1e-6)
    cos_angle = torch.sum(view_dir * normal, dim=1)
    cos_mask = cos_angle < 0.0
    normal[cos_mask] = -normal[cos_mask]

    keep_mask = torch.abs(cos_angle) > np.cos(np.deg2rad(85.0))
    xyz, normal = xyz[keep_mask], normal[keep_mask]

    return xyz, normal, None
//  [python]
def normal_func(xyz: torch.Tensor, normal: torch.Tensor, sensor: torch.Tensor):
    assert normal is None, "normal already exists"
    assert sensor is not None, "please provide sensor positions for consistent orientations"

    xyz_numpy = xyz.cpu().numpy()
    indices, normal = pcu.estimate_point_cloud_normals_knn(xyz_numpy, 64)
    normal = torch.from_numpy(normal).to(xyz)
    indices = torch.from_numpy(indices).to(xyz).long()

    xyz, sensor = xyz[indices], sensor[indices]

    view_dir = sensor - xyz
    view_dir = view_dir / (torch.linalg.norm(view_dir, dim=-1, keepdim=True) + 1e-6)
    cos_angle = torch.sum(view_dir * normal, dim=1)
    cos_mask = cos_angle < 0.0
    normal[cos_mask] = -normal[cos_mask]

    keep_mask = torch.abs(cos_angle) > np.cos(np.deg2rad(85.0))
    xyz, normal = xyz[keep_mask], normal[keep_mask]

    return xyz, normal, None
//  [python]
def retraction(X, G, A, lr, method='SVD', adapt=False):
    r"""retract the updated tensor onto Stiefel manifold
    Args:
        X (tensor): tensor of parameters in matrix form
        G (tensor): gradient tensor in matrix form
        A (tensor): the tensor G = AX in matrix form
        lr (float): learning rate
        method (str, optional): method for retraction into Stiefel manifold ('SVD', 'Cayley') (default: 'SVD')
    """
    if adapt:
        lr = min(lr, 4.0 / (torch.norm(A).item() + 1e-8)) 
    if method == 'SVD':
        U, _, Vt = LA.svd((X - lr * G).cpu().numpy(), full_matrices=False)
        X_out = torch.tensor(U @ Vt)
        # U, _, V = svd_(X - lr * G)
        # X_out = U @ V.t()

    elif method == 'Cayley':
        X_out = X - lr * G
        for i in range(40):
            X_out = X - 0.3 * lr * A @ (X + X_out)

    return X_out
//  [python]
def retraction(X, G, A, lr, method='SVD', adapt=False):
    r"""retract the updated tensor onto Stiefel manifold
    Args:
        X (tensor): tensor of parameters in matrix form
        G (tensor): gradient tensor in matrix form
        A (tensor): the tensor G = AX in matrix form
        lr (float): learning rate
        method (str, optional): method for retraction into Stiefel manifold ('SVD', 'Cayley') (default: 'SVD')
    """
    if adapt:
        lr = min(lr, 4.0 / (torch.linalg.norm(A).item() + 1e-8)) 
    if method == 'SVD':
        U, _, Vt = LA.svd((X - lr * G).cpu().numpy(), full_matrices=False)
        X_out = torch.tensor(U @ Vt)
        # U, _, V = svd_(X - lr * G)
        # X_out = U @ V.t()

    elif method == 'Cayley':
        X_out = X - lr * G
        for i in range(40):
            X_out = X - 0.3 * lr * A @ (X + X_out)

    return X_out
//  [python]
def quaternion_to_axis_angle(q):
  angle = 2 * torch.acos(torch.clip(q[:, 3], -0.99999, 0.99999))[:, None]
  norm = torch.clip(torch.norm(q[:, :3], dim=1), 1e-5, 1)[:, None]
  axis = q[:, :3] / norm
  return axis, angle
//  [python]
def quaternion_to_axis_angle(q):
  angle = 2 * torch.acos(torch.clip(q[:, 3], -0.99999, 0.99999))[:, None]
  norm = torch.clip(torch.linalg.norm(q[:, :3], dim=1), 1e-5, 1)[:, None]
  axis = q[:, :3] / norm
  return axis, angle
//  [python]
def trans_ang_loss(t, tgt):
    """Computes L1 loss for translation vector ANGULAR error
    Input:
    t - estimated translation vector [B, 1, 3]
    tgt - ground-truth translation vector [B, 1, 3]
    Output: translation_loss
    """

    scale_t = torch.norm(t, dim=-1)
    scale_tgt = torch.norm(tgt, dim=-1)

    cosine = (t @ tgt.transpose(1, 2)).squeeze(-1) / (scale_t * scale_tgt + 1e-6)
    cosine = torch.clip(cosine, -0.99999, 0.99999)  # handle numerical errors and NaNs
    t_ang_err = torch.acos(cosine)
    t_ang_err = torch.minimum(t_ang_err, math.pi - t_ang_err)
    return torch.abs(t_ang_err - torch.zeros_like(t_ang_err)), t_ang_err
//  [python]
def trans_ang_loss(t, tgt):
    """Computes L1 loss for translation vector ANGULAR error
    Input:
    t - estimated translation vector [B, 1, 3]
    tgt - ground-truth translation vector [B, 1, 3]
    Output: translation_loss
    """

    scale_t = torch.linalg.norm(t, dim=-1)
    scale_tgt = torch.linalg.norm(tgt, dim=-1)

    cosine = (t @ tgt.transpose(1, 2)).squeeze(-1) / (scale_t * scale_tgt + 1e-6)
    cosine = torch.clip(cosine, -0.99999, 0.99999)  # handle numerical errors and NaNs
    t_ang_err = torch.acos(cosine)
    t_ang_err = torch.minimum(t_ang_err, math.pi - t_ang_err)
    return torch.abs(t_ang_err - torch.zeros_like(t_ang_err)), t_ang_err
//  [python]
def norm_tensor(tensor: torch.Tensor) -> torch.Tensor:
    return tensor / tensor.norm(p=2, dim=-1, keepdim=True)
//  [python]
def norm_tensor(tensor: torch.Tensor) -> torch.Tensor:
    return tensor / torch.linalg.norm(tensor, ord=2, dim=-1, keepdim=True)
//  [python]
def forward(self, states):
        N, state_dim = states.shape
        assert self.state_dim == state_dim
        pos = states[:, self.idx_pos_mask]
        bias = states[:, self.idx_b_mask]
        expected_observation = torch.norm(self.satXYZb[None, :, :3] - pos[:, None, :], dim=-1) + torch.abs(self.satXYZb[None, :, 3] - bias)
        R = torch.eye(self.observation_dim).expand((N, self.observation_dim, self.observation_dim)) * self.pr_std
        
        return expected_observation, R
//  [python]
def forward(self, states):
        N, state_dim = states.shape
        assert self.state_dim == state_dim
        pos = states[:, self.idx_pos_mask]
        bias = states[:, self.idx_b_mask]
        expected_observation = torch.linalg.norm(self.satXYZb[None, :, :3] - pos[:, None, :], dim=-1) + torch.abs(self.satXYZb[None, :, 3] - bias)
        R = torch.eye(self.observation_dim).expand((N, self.observation_dim, self.observation_dim)) * self.pr_std
        
        return expected_observation, R
//  [python]
def rmsprop(params: List[Tensor],
            grads: List[Tensor],
            square_avgs: List[Tensor],
            grad_avgs: List[Tensor],
            momentum_buffer_list: List[Tensor],
            lr: float,
            alpha: float,
            eps: float,
            weight_decay: float,
            momentum: float,
            centered: bool,
            method:str):
    r"""Functional API that performs rmsprop algorithm computation.
    """
    keep_dir = True

    for i, param in enumerate(params):
        grad = grads[i]
        square_avg = square_avgs[i]
        buf = momentum_buffer_list[i]

        if weight_decay != 0:
            grad = grad.add(param, alpha=weight_decay)

        square_avg.mul_(alpha).addcmul_(grad, grad, value=1 - alpha)

        if centered:
            grad_avg = grad_avgs[i]
            grad_avg.mul_(alpha).add_(grad, alpha=1 - alpha)
            avg = square_avg.addcmul(grad_avg, grad_avg, value=-0.5).sqrt_().add_(eps)
        else:
            avg = square_avg.sqrt().add_(eps)

        # buf.mul_(momentum).addcdiv_(grad, avg)
        # param.add_(buf, alpha=-lr)


        r = torch.norm(avg)
        if keep_dir:
            buf = momentum * buf + grad / r 
        else: 
            buf = momentum * buf + grad

        d_p = buf

        dims = param.size()
        bond_in, _ = param.leg_type
        dim1 = np.prod(dims[:bond_in])
        dim2 = np.prod(dims[bond_in:])
        X = param.data.view(dim1, dim2)
        dX = d_p.data.view(dim1, dim2)

        G, A = grad_proj(X, dX)
        momentum_buffer_list[i] = torch.clone(G.view(dims)).detach()

        if not keep_dir:
            G = torch.div(G, avg.view(dim1, dim2)) / 4
            G, A = grad_proj(X, G)

        X_out = retraction(X, G, A / r, lr, method=method, adapt=False)
        param.data.fill_(0)
        param.data += X_out.view(dims).to(param.device)
//  [python]
def rmsprop(params: List[Tensor],
            grads: List[Tensor],
            square_avgs: List[Tensor],
            grad_avgs: List[Tensor],
            momentum_buffer_list: List[Tensor],
            lr: float,
            alpha: float,
            eps: float,
            weight_decay: float,
            momentum: float,
            centered: bool,
            method:str):
    r"""Functional API that performs rmsprop algorithm computation.
    """
    keep_dir = True

    for i, param in enumerate(params):
        grad = grads[i]
        square_avg = square_avgs[i]
        buf = momentum_buffer_list[i]

        if weight_decay != 0:
            grad = grad.add(param, alpha=weight_decay)

        square_avg.mul_(alpha).addcmul_(grad, grad, value=1 - alpha)

        if centered:
            grad_avg = grad_avgs[i]
            grad_avg.mul_(alpha).add_(grad, alpha=1 - alpha)
            avg = square_avg.addcmul(grad_avg, grad_avg, value=-0.5).sqrt_().add_(eps)
        else:
            avg = square_avg.sqrt().add_(eps)

        # buf.mul_(momentum).addcdiv_(grad, avg)
        # param.add_(buf, alpha=-lr)


        r = torch.linalg.norm(avg)
        if keep_dir:
            buf = momentum * buf + grad / r 
        else: 
            buf = momentum * buf + grad

        d_p = buf

        dims = param.size()
        bond_in, _ = param.leg_type
        dim1 = np.prod(dims[:bond_in])
        dim2 = np.prod(dims[bond_in:])
        X = param.data.view(dim1, dim2)
        dX = d_p.data.view(dim1, dim2)

        G, A = grad_proj(X, dX)
        momentum_buffer_list[i] = torch.clone(G.view(dims)).detach()

        if not keep_dir:
            G = torch.div(G, avg.view(dim1, dim2)) / 4
            G, A = grad_proj(X, G)

        X_out = retraction(X, G, A / r, lr, method=method, adapt=False)
        param.data.fill_(0)
        param.data += X_out.view(dims).to(param.device)
//  [python]
def ps_gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=int(0.00*self.W), high=int(1.00*self.W), size=[batch_size])
        pixels_y = torch.randint(low=int(0.00*self.H), high=int(1.00*self.H), size=[batch_size])
        color = self.images[img_idx[0]][img_idx[1]][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx[0]][(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[0], 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[0], None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[img_idx[0], 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()
//  [python]
def ps_gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=int(0.00*self.W), high=int(1.00*self.W), size=[batch_size])
        pixels_y = torch.randint(low=int(0.00*self.H), high=int(1.00*self.H), size=[batch_size])
        color = self.images[img_idx[0]][img_idx[1]][(pixels_y, pixels_x)]    # batch_size, 3
        mask = self.masks[img_idx[0]][(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[0], 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[0], None, :3, :3], rays_v[:, :, None]).squeeze()  # batch_size, 3
        rays_o = self.pose_all[img_idx[0], 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()
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=int(0.00*self.W), high=int(1.00*self.W), size=[batch_size])
        pixels_y = torch.randint(low=int(0.00*self.H), high=int(1.00*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()
//  [python]
def gen_random_rays_at(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=int(0.00*self.W), high=int(1.00*self.W), size=[batch_size])
        pixels_y = torch.randint(low=int(0.00*self.H), high=int(1.00*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()
//  [python]
def cholesky(self, quat):
        # Return the covariance matrix of the noise.
        with torch.enable_grad():
            tmp = quat.detach().clone()
            N, D = tmp.shape
            tmp = tmp[:, None, :].expand((N, 3, D))
            tmp.requires_grad = True
            eul = quat2eul(tmp.reshape(-1, D)).reshape(N, -1, 3)
            mask = torch.eye(3, device=eul.device).repeat(N, 1, 1)
            quat_jacobian = torch.autograd.grad(eul, tmp, mask, create_graph=False)[0].detach().float()
            quat_jacobian = quat_jacobian.transpose(1, 2)
            ret_chol = torch.matmul(quat_jacobian, torch.diag(torch.stack([self.r_std, self.p_std, self.y_std])).float())
            ret_chol = torch.norm(ret_chol, dim=2)
            ret_chol = torch.diag_embed(ret_chol)
        return ret_chol
//  [python]
def cholesky(self, quat):
        # Return the covariance matrix of the noise.
        with torch.enable_grad():
            tmp = quat.detach().clone()
            N, D = tmp.shape
            tmp = tmp[:, None, :].expand((N, 3, D))
            tmp.requires_grad = True
            eul = quat2eul(tmp.reshape(-1, D)).reshape(N, -1, 3)
            mask = torch.eye(3, device=eul.device).repeat(N, 1, 1)
            quat_jacobian = torch.autograd.grad(eul, tmp, mask, create_graph=False)[0].detach().float()
            quat_jacobian = quat_jacobian.transpose(1, 2)
            ret_chol = torch.matmul(quat_jacobian, torch.diag(torch.stack([self.r_std, self.p_std, self.y_std])).float())
            ret_chol = torch.linalg.norm(ret_chol, dim=2)
            ret_chol = torch.diag_embed(ret_chol)
        return ret_chol
//  [python]
def eigsh_py(A: Tensor,
             largest: bool = True,
             m: int = 20,
             max_iter: int = 10000,
             tol: float = 1e-5,
             seed: Optional[int] = None
             ) -> Tuple[Tensor, Tensor, int]:
    """Pure-python variant of eigsh"""

    if seed is not None:
        torch.manual_seed(seed)

    n, _ = A.shape
    m = min(n, m)
    dtype = A.dtype
    device = A.device

    # lanczos buffers
    V = torch.zeros(m, n, dtype=dtype, device=device)
    T = torch.zeros(m, m, dtype=dtype, device=device)
    beta = torch.tensor(0., dtype=dtype, device=device)
    r = torch.empty(n, dtype=dtype, device=device)

    # eigh buffers
    eigvals = torch.zeros(m, dtype=dtype, device=device)
    eigvecs = torch.zeros(m, m, dtype=dtype, device=device).t()

    # ritz buffers
    w = torch.zeros(1, dtype=dtype, device=device)
    z = torch.zeros(m, dtype=dtype, device=device)

    # begin restarted Lanczos iterations
    n_iter = 0
    V[0].normal_()
    while n_iter <= max_iter:
        n_iter += 1

        # Perform m-step Lanczos factorization starting at V[0]
        V[0].div_(torch.norm(V[0]))
        mk = m
        for i in range(m):
            torch.matmul(A, V[i], out=r)
            if i > 0:
                r.sub_(beta * V[i-1])
            torch.dot(V[i], r, out=T[i, i])
            r.sub_(T[i, i] * V[i])
            beta = torch.norm(r)
            if beta < 1e-10:  # TODO: avoid GPU synchronization
                mk = i + 1
                break
            if i + 1 < m:
                T[i, i+1] = beta
                T[i+1, i] = beta
                torch.div(r, beta, out=V[i+1])

        if mk == 1:
            # Lanczos factorization exited after 1 step.
            # We have quick access to eigval/vec so we will terminate now.
            w[0] = T[0, 0]
            break
        elif mk != m:
            # Lanczos factorization exited early.
            # We will compute eigval/vec and exit.
            V = V[:mk]
            T = T[:mk, :mk].contiguous()
            eigvals = eigvals[:mk]
            eigvecs = eigvecs[:mk, :mk].contiguous().t()
            z = z[:mk]

        # Compute largest/smallest eigenpair of T
        # TODO: this computation step could be made much more efficient
        #  by the following:
        #     1. Capitalize on the tridiagonal structure of T
        #     2. Only compute the eigenpairs we need
        if torch.jit.is_scripting():
            torch.linalg.eigh(T, eigvals=eigvals, eigvecs=eigvecs)
        else:
            torch.linalg.eigh(T, out=(eigvals, eigvecs))

        # locate eigval and compute projection
        if largest:
            w[0] = eigvals[-1]
            z[:] = eigvecs[:, -1]
        else:
            w[0] = eigvals[0]
            z[:] = eigvecs[:, 0]

        # compute projection of eigvec back to R^n
        V[0] = torch.matmul(z, V)

        if mk != m or beta * z[-1].abs() < tol:  # TODO: avoid GPU synchronization
            break

    return w, V[0], n_iter
//  [python]
def eigsh_py(A: Tensor,
             largest: bool = True,
             m: int = 20,
             max_iter: int = 10000,
             tol: float = 1e-5,
             seed: Optional[int] = None
             ) -> Tuple[Tensor, Tensor, int]:
    """Pure-python variant of eigsh"""

    if seed is not None:
        torch.manual_seed(seed)

    n, _ = A.shape
    m = min(n, m)
    dtype = A.dtype
    device = A.device

    # lanczos buffers
    V = torch.zeros(m, n, dtype=dtype, device=device)
    T = torch.zeros(m, m, dtype=dtype, device=device)
    beta = torch.tensor(0., dtype=dtype, device=device)
    r = torch.empty(n, dtype=dtype, device=device)

    # eigh buffers
    eigvals = torch.zeros(m, dtype=dtype, device=device)
    eigvecs = torch.zeros(m, m, dtype=dtype, device=device).t()

    # ritz buffers
    w = torch.zeros(1, dtype=dtype, device=device)
    z = torch.zeros(m, dtype=dtype, device=device)

    # begin restarted Lanczos iterations
    n_iter = 0
    V[0].normal_()
    while n_iter <= max_iter:
        n_iter += 1

        # Perform m-step Lanczos factorization starting at V[0]
        V[0].div_(torch.linalg.norm(V[0]))
        mk = m
        for i in range(m):
            torch.matmul(A, V[i], out=r)
            if i > 0:
                r.sub_(beta * V[i-1])
            torch.dot(V[i], r, out=T[i, i])
            r.sub_(T[i, i] * V[i])
            torch.linalg.norm(r, out=beta)
            if beta < 1e-10:  # TODO: avoid GPU synchronization
                mk = i + 1
                break
            if i + 1 < m:
                T[i, i+1] = beta
                T[i+1, i] = beta
                torch.div(r, beta, out=V[i+1])

        if mk == 1:
            # Lanczos factorization exited after 1 step.
            # We have quick access to eigval/vec so we will terminate now.
            w[0] = T[0, 0]
            break
        elif mk != m:
            # Lanczos factorization exited early.
            # We will compute eigval/vec and exit.
            V = V[:mk]
            T = T[:mk, :mk].contiguous()
            eigvals = eigvals[:mk]
            eigvecs = eigvecs[:mk, :mk].contiguous().t()
            z = z[:mk]

        # Compute largest/smallest eigenpair of T
        # TODO: this computation step could be made much more efficient
        #  by the following:
        #     1. Capitalize on the tridiagonal structure of T
        #     2. Only compute the eigenpairs we need
        if torch.jit.is_scripting():
            torch.linalg.eigh(T, eigvals=eigvals, eigvecs=eigvecs)
        else:
            torch.linalg.eigh(T, out=(eigvals, eigvecs))

        # locate eigval and compute projection
        if largest:
            w[0] = eigvals[-1]
            z[:] = eigvecs[:, -1]
        else:
            w[0] = eigvals[0]
            z[:] = eigvecs[:, 0]

        # compute projection of eigvec back to R^n
        V[0] = torch.matmul(z, V)

        if mk != m or beta * z[-1].abs() < tol:  # TODO: avoid GPU synchronization
            break

    return w, V[0], n_iter
//  [python]
def ps_gen_random_rays_at_view_on_all_lights(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=int(0.00*self.W), high=int(1.00*self.W), size=[batch_size], device='cpu')
        pixels_y = torch.randint(low=int(0.00*self.H), high=int(1.00*self.H), size=[batch_size], device='cpu')
        images_warmup = self.images_warmup[img_idx,:,pixels_y,pixels_x,:] # nb_light, batch_size, 3-4
        images = self.images[img_idx,:,pixels_y,pixels_x,:] # nb_light, batch_size, 3-4

        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().cuda()  # 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(), mask[:, :1].cpu()], dim=-1).cuda(), images_warmup.cuda(), images.cuda(), pixels_x.cuda(), pixels_y.cuda()
//  [python]
def ps_gen_random_rays_at_view_on_all_lights(self, img_idx, batch_size):
        """
        Generate random rays at world space from one camera.
        """
        pixels_x = torch.randint(low=int(0.00*self.W), high=int(1.00*self.W), size=[batch_size], device='cpu')
        pixels_y = torch.randint(low=int(0.00*self.H), high=int(1.00*self.H), size=[batch_size], device='cpu')
        images_warmup = self.images_warmup[img_idx,:,pixels_y,pixels_x,:] # nb_light, batch_size, 3-4
        images = self.images[img_idx,:,pixels_y,pixels_x,:] # nb_light, batch_size, 3-4

        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().cuda()  # 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(), mask[:, :1].cpu()], dim=-1).cuda(), images_warmup.cuda(), images.cuda(), pixels_x.cuda(), pixels_y.cuda()
//  [python]
def broyden(g, x_init, J_inv_init, max_steps=10, 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 step in range(max_steps):
        # update parameter 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 future 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 x_opt, gx_norm_opt, gx_norm_opt < cvg_thresh
//  [python]
def broyden(g, x_init, J_inv_init, max_steps=10, 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 step in range(max_steps):
        # update parameter 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 future 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 x_opt, gx_norm_opt, gx_norm_opt < cvg_thresh
//  [python]
def cholesky(self, vel_enu, orientation, delta_quat):
        # Return the covariance matrix of the noise.
        N, _ = vel_enu.shape
        orientation = orientation[:, None, :].expand((N, self.observation_dim, 4)).reshape(-1, 4)
        rmat = rmat[:, None, :, :].expand((N, self.observation_dim, 3, 3)).reshape(-1, 3, 3)
        with torch.enable_grad():
            tmp = vel_enu.detach().clone()
            tmp = tmp[:, None, :].expand((N, self.observation_dim, 3))
            tmp.requires_grad = True
            img_pts = self._vel_to_2dfeatures(tmp.reshape(-1, 3), orientation, rmat).reshape(N, -1, self.observation_dim)
            mask = torch.eye(self.observation_dim, device=img_pts.device).repeat(N, 1, 1)
            feat_jacobian = torch.autograd.grad(img_pts, tmp, mask, create_graph=False)[0].detach().float()
            ret_chol = torch.matmul(feat_jacobian, self.prev_covariance[0, :3, :3])
            ret_chol = torch.norm(ret_chol, dim=2)
            ret_chol = torch.diag_embed(ret_chol)
        return ret_chol
//  [python]
def cholesky(self, vel_enu, orientation, delta_quat):
        # Return the covariance matrix of the noise.
        N, _ = vel_enu.shape
        orientation = orientation[:, None, :].expand((N, self.observation_dim, 4)).reshape(-1, 4)
        rmat = rmat[:, None, :, :].expand((N, self.observation_dim, 3, 3)).reshape(-1, 3, 3)
        with torch.enable_grad():
            tmp = vel_enu.detach().clone()
            tmp = tmp[:, None, :].expand((N, self.observation_dim, 3))
            tmp.requires_grad = True
            img_pts = self._vel_to_2dfeatures(tmp.reshape(-1, 3), orientation, rmat).reshape(N, -1, self.observation_dim)
            mask = torch.eye(self.observation_dim, device=img_pts.device).repeat(N, 1, 1)
            feat_jacobian = torch.autograd.grad(img_pts, tmp, mask, create_graph=False)[0].detach().float()
            ret_chol = torch.matmul(feat_jacobian, self.prev_covariance[0, :3, :3])
            ret_chol = torch.linalg.norm(ret_chol, dim=2)
            ret_chol = torch.diag_embed(ret_chol)
        return ret_chol
//  [python]
def embedding_clip(param, R: int):
    """Clip embedding vector according to $R$

    Arguments:
        param (tensor): The embedding vector
        R (int): The upper bound of ratings
    """
    # Turn all negative entries of U into 0
    param.data = (torch.abs(param.data) + param.data) * 0.5
    # Clip tensor
    norms = torch.norm(param.data, dim=1)
    threshold = np.sqrt(R)
    param.data[norms > threshold] *= (threshold /
                                      norms[norms > threshold]).reshape(
                                          (-1, 1))
    param.data[param.data < 0] = 0.
//  [python]
def embedding_clip(param, R: int):
    """Clip embedding vector according to $R$

    Arguments:
        param (tensor): The embedding vector
        R (int): The upper bound of ratings
    """
    # Turn all negative entries of U into 0
    param.data = (torch.abs(param.data) + param.data) * 0.5
    # Clip tensor
    norms = torch.linalg.norm(param.data, dim=1)
    threshold = np.sqrt(R)
    param.data[norms > threshold] *= (threshold /
                                      norms[norms > threshold]).reshape(
                                          (-1, 1))
    param.data[param.data < 0] = 0.
//  [python]
def randomized_power_err_est(A, U, Lambda_hat, q, dtype, device):
    g = torch.randn(A.shape[0], dtype=dtype, device=device)
    v_0 = g / torch.norm(g)
    E_hat = torch.inf
    for _ in range(q):
        v = A @ v_0 - U @ (Lambda_hat * (U.T @ v_0))
        E_hat = (v_0 * v).sum()
        v = v / torch.norm(v)
        v_0 = v
    return E_hat
//  [python]
def randomized_power_err_est(A, U, Lambda_hat, q, dtype, device):
    g = torch.randn(A.shape[0], dtype=dtype, device=device)
    v_0 = g / torch.linalg.norm(g)
    E_hat = torch.inf
    for _ in range(q):
        v = A @ v_0 - U @ (Lambda_hat * (U.T @ v_0))
        E_hat = (v_0 * v).sum()
        v = v / torch.linalg.norm(v)
        v_0 = v
    return E_hat
//  [python]
def update_model( self ):
        dataset_train = self.build_dataset()
        self._debug_mode = False
        if self._debug_mode and os.path.exists(self._feature_cache):
            print("using feature cache")
            with open(self._feature_cache, 'rb') as f:
                train_data = torch.load(f)
        else:
            print("generating bbox features")
            train_data = self.get_features(dataset_train)
            saved = torch.save(train_data, self._feature_cache)
        train_data = torch.norm(train_data, dim=1, p=self._norm_degree)
        
        self.remax_model = ReMax(train_data)

        self.save_final_model()