jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_norm_complex(self, device, dtype):
        def gen_error_message(input_size, ord, keepdim, dim=None):
            return "complex norm failed for input size %s, ord=%s, keepdim=%s, dim=%s" % (
                input_size, ord, keepdim, dim)

        vector_ords = [None, 0, 1, 2, 3, inf, -1, -2, -3, -inf]
        matrix_ords = [None, 'fro', 'nuc', 1, 2, inf, -1, -2, -inf]

        # Test supported ords
        for keepdim in [False, True]:
            # vector norm
            x = torch.randn(25, device=device, dtype=dtype)
            xn = x.cpu().numpy()
            for ord in vector_ords:
                res = torch.linalg.norm(x, ord, keepdim=keepdim).cpu()
                expected = np.linalg.norm(xn, ord, keepdims=keepdim)
                msg = gen_error_message(x.size(), ord, keepdim)
                self.assertEqual(res.shape, expected.shape, msg=msg)
                self.assertEqual(res, expected, msg=msg)

                res_out = torch.tensor([]).to(device)
                torch.linalg.norm(x, ord, keepdim=keepdim, out=res_out)
                self.assertEqual(res_out.shape, expected.shape, msg=msg)
                self.assertEqual(res_out.cpu(), expected, msg=msg)

            # matrix norm
            x = torch.randn(25, 25, device=device, dtype=dtype)
            xn = x.cpu().numpy()
            for ord in matrix_ords:
                res = torch.linalg.norm(x, ord, keepdim=keepdim).cpu()
                expected = np.linalg.norm(xn, ord, keepdims=keepdim)
                msg = gen_error_message(x.size(), ord, keepdim)
                self.assertEqual(res.shape, expected.shape, msg=msg)
                self.assertEqual(res, expected, msg=msg)

                res_out = torch.tensor([]).to(device)
                torch.linalg.norm(x, ord, keepdim=keepdim, out=res_out)
                self.assertEqual(res_out.shape, expected.shape, msg=msg)
                self.assertEqual(res_out.cpu(), expected, msg=msg)
//  [python]
def forward(self, text):
        tokens = tokenize_with_truncation(text, truncate=True).to(self.device)
        z = self.model.encode_text(tokens)
        if self.normalize:
            z = z / torch.norm(z, dim=1, keepdim=True)
        return z
//  [python]
def forward(self, text):
        # tokens = clip.tokenize(text).to(self.device)
        tokens = tokenize_with_truncation(text, truncate=True).to(self.device)
        z = self.model.encode_text(tokens)
        if self.normalize:
            z = z / torch.linalg.norm(z, dim=1, keepdim=True)
        return z
//  [python]
def forward(self, text):
        tok_text = self.tokenizer(text).to(self.device)
        z = self.model.encode_text(tok_text)
        if self.normalize:
            z = z / torch.norm(z, dim=1, keepdim=True)
        return z
//  [python]
def forward(self, text):
        tok_text = self.tokenizer(text).to(self.device)
        z = self.model.encode_text(tok_text)
        if self.normalize:
            z = z / torch.linalg.norm(z, dim=1, keepdim=True)
        return z
//  [python]
def build_all_rays(H, W, fx, fy, cx, cy, c2w, device, nerf_coordinate=True, dir_normalize=False):
    """
    Build rays for a whole image.
    """
    # !!! do not normalize the dir, otherwise the scale of pose will be deteriorated
    dir_normalize = DIR_NORMALIZE
    if isinstance(c2w, np.ndarray):
        c2w = torch.from_numpy(c2w).to(device)

    x, y = torch.meshgrid(
        torch.linspace(0, W-1, W).to(device),
        torch.linspace(0, H-1, H).to(device),
        indexing='ij',
    )
    x, y = x.t(), y.t()  # [W, H] -> [H, W]
    if nerf_coordinate:
        # https://github.com/cvg/nice-slam/issues/10#issuecomment-1128590013
        dirs = torch.stack([
            (x - cx) / fx,
            -(y - cy) / fy,
            -torch.ones_like(x),
            ], dim=-1).to(device)  # (H, W, 3)
    else:
        dirs = torch.stack([
            (x - cx) / fx,
            (y - cy) / fy,
            torch.ones_like(x),
        ], dim=-1).to(device)  # (H, W, 3)

    if dir_normalize:
        dirs = dirs / torch.norm(dirs, dim=-1, keepdim=True)  # H, W, 3
        # comment it if you are sure the pose is fine
        raise TypeError("Ray direction shouldn't be normalized, otherwise the scale of pose will be destroyed!")

    rays_d = dirs @ c2w[:3, :3].t()  # [H, W, 3]
    rays_o = c2w[:3, 3].reshape(1, 1, 3).repeat(H, W, 1)  # [n_rays, 3]

    return rays_o, rays_d
//  [python]
def build_all_rays(H, W, fx, fy, cx, cy, c2w, device, nerf_coordinate=True, dir_normalize=False):
    """
    Build rays for a whole image.
    """
    # !!! do not normalize the dir, otherwise the scale of pose will be deteriorated
    dir_normalize = DIR_NORMALIZE
    if isinstance(c2w, np.ndarray):
        c2w = torch.from_numpy(c2w).to(device)

    x, y = torch.meshgrid(
        torch.linspace(0, W-1, W).to(device),
        torch.linspace(0, H-1, H).to(device),
        indexing='ij',
    )
    x, y = x.t(), y.t()  # [W, H] -> [H, W]
    if nerf_coordinate:
        # https://github.com/cvg/nice-slam/issues/10#issuecomment-1128590013
        dirs = torch.stack([
            (x - cx) / fx,
            -(y - cy) / fy,
            -torch.ones_like(x),
            ], dim=-1).to(device)  # (H, W, 3)
    else:
        dirs = torch.stack([
            (x - cx) / fx,
            (y - cy) / fy,
            torch.ones_like(x),
        ], dim=-1).to(device)  # (H, W, 3)

    if dir_normalize:
        dirs = dirs / torch.linalg.norm(dirs, ord=2, dim=-1, keepdim=True)  # H, W, 3
        # comment it if you are sure the pose is fine
        raise TypeError("Ray direction shouldn't be normalized, otherwise the scale of pose will be destroyed!")

    rays_d = dirs @ c2w[:3, :3].t()  # [H, W, 3]
    rays_o = c2w[:3, 3].reshape(1, 1, 3).repeat(H, W, 1)  # [n_rays, 3]

    return rays_o, rays_d
//  [python]
def cam_pts_2_angle(self, cam_pts):
        x = cam_pts[:, 0]
        y = cam_pts[:, 1]
        z = cam_pts[:, 2]
        distance = torch.norm(cam_pts, p=2, dim=1)
        v_angle = torch.acos(-y/distance)/ math.pi * 180 # wrt minus y direction
        h_angle = 180 - torch.atan2(z, x)/ math.pi * 180 # wrt to x direction
        return v_angle, h_angle, distance
//  [python]
def cam_pts_2_angle(self, cam_pts):
        x = cam_pts[:, 0]
        y = cam_pts[:, 1]
        z = cam_pts[:, 2]
        distance = torch.linalg.norm(cam_pts, ord=2, dim=1)
        v_angle = torch.acos(-y/distance)/ math.pi * 180 # wrt minus y direction
        h_angle = 180 - torch.atan2(z, x)/ math.pi * 180 # wrt to x direction
        return v_angle, h_angle, distance
//  [python]
def refine_ray_generation(campos,
                               raydir,
                               point_count,
                               prev_ts,
                               prev_weights,
                               domain_size=1.,
                               jitter=0,
                               **kargs):
    # inputs
    # campos: N x 3
    # raydir: N x Rays x 3, must be normalized
    # point_count: int
    # prev_ts: N x Rays x PrevSamples
    # prev_weights: N x Rays x PrevSamples
    # outputs
    # raypos: N x Rays x Samples x 3
    # segment_length: N x Rays x Samples
    # valid: N x Rays x Samples
    # ts: N x Rays x Samples
    with torch.no_grad():
        end_point_ts = sample_pdf(prev_ts[..., None], prev_weights,
                                  point_count + 1, jitter <= 0)
        end_point_ts = end_point_ts.view(end_point_ts.shape[:-1])
        segment_length = end_point_ts[:, :, 1:] - end_point_ts[:, :, :-1]
        middle_point_ts = (end_point_ts[:, :, :-1] +
                           end_point_ts[:, :, 1:]) / 2
        raypos = campos[:, None,
                        None, :] + raydir[:, :,
                                          None, :] * middle_point_ts[:, :, :,
                                                                     None]
        valid = torch.ones_like(middle_point_ts,
                                dtype=middle_point_ts.dtype,
                                device=middle_point_ts.device)
    segment_length*=torch.norm(raydir[..., None, :], dim=-1)
    return raypos, segment_length, valid, middle_point_ts
//  [python]
def refine_ray_generation(campos,
                               raydir,
                               point_count,
                               prev_ts,
                               prev_weights,
                               domain_size=1.,
                               jitter=0,
                               **kargs):
    # inputs
    # campos: N x 3
    # raydir: N x Rays x 3, must be normalized
    # point_count: int
    # prev_ts: N x Rays x PrevSamples
    # prev_weights: N x Rays x PrevSamples
    # outputs
    # raypos: N x Rays x Samples x 3
    # segment_length: N x Rays x Samples
    # valid: N x Rays x Samples
    # ts: N x Rays x Samples
    with torch.no_grad():
        end_point_ts = sample_pdf(prev_ts[..., None], prev_weights,
                                  point_count + 1, jitter <= 0)
        end_point_ts = end_point_ts.view(end_point_ts.shape[:-1])
        segment_length = end_point_ts[:, :, 1:] - end_point_ts[:, :, :-1]
        middle_point_ts = (end_point_ts[:, :, :-1] +
                           end_point_ts[:, :, 1:]) / 2
        raypos = campos[:, None,
                        None, :] + raydir[:, :,
                                          None, :] * middle_point_ts[:, :, :,
                                                                     None]
        valid = torch.ones_like(middle_point_ts,
                                dtype=middle_point_ts.dtype,
                                device=middle_point_ts.device)
    segment_length*=torch.linalg.norm(raydir[..., None, :], axis=-1)
    return raypos, segment_length, valid, middle_point_ts
//  [python]
def nerf_refine_ray_generation(campos,
                               raydir,
                               point_count,
                               prev_ts,
                               prev_weights,
                               domain_size=1.,
                               jitter=0,
                               **kargs):
    # inputs
    # campos: N x 3
    # raydir: N x Rays x 3, must be normalized
    # point_count: int
    # prev_ts: N x Rays x PrevSamples,  uniformed depth segments between near and far
    # prev_weights: N x Rays x PrevSamples
    # outputs
    # raypos: N x Rays x Samples x 3
    # segment_length: N x Rays x Samples
    # valid: N x Rays x Samples
    # ts: N x Rays x Samples
    with torch.no_grad():

        end_point_ts = sample_pdf(prev_ts[..., None], prev_weights,
                                  point_count + 1, jitter <= 0)
        end_point_ts = end_point_ts.view(end_point_ts.shape[:-1])
        segment_length = end_point_ts[:, :, 1:] - end_point_ts[:, :, :-1]
        segment_length *= torch.norm(raydir[..., None, :], dim=-1)
        middle_point_ts = (end_point_ts[:, :, :-1] +
                           end_point_ts[:, :, 1:]) / 2
        raypos = campos[:, None,
                        None, :] + raydir[:, :,
                                          None, :] * middle_point_ts[:, :, :,
                                                                     None]
        valid = torch.ones_like(middle_point_ts,
                                dtype=middle_point_ts.dtype,
                                device=middle_point_ts.device)

    return raypos, segment_length, valid, middle_point_ts
//  [python]
def nerf_refine_ray_generation(campos,
                               raydir,
                               point_count,
                               prev_ts,
                               prev_weights,
                               domain_size=1.,
                               jitter=0,
                               **kargs):
    # inputs
    # campos: N x 3
    # raydir: N x Rays x 3, must be normalized
    # point_count: int
    # prev_ts: N x Rays x PrevSamples,  uniformed depth segments between near and far
    # prev_weights: N x Rays x PrevSamples
    # outputs
    # raypos: N x Rays x Samples x 3
    # segment_length: N x Rays x Samples
    # valid: N x Rays x Samples
    # ts: N x Rays x Samples
    with torch.no_grad():

        end_point_ts = sample_pdf(prev_ts[..., None], prev_weights,
                                  point_count + 1, jitter <= 0)
        end_point_ts = end_point_ts.view(end_point_ts.shape[:-1])
        segment_length = end_point_ts[:, :, 1:] - end_point_ts[:, :, :-1]
        segment_length *= torch.linalg.norm(raydir[..., None, :], axis=-1)
        middle_point_ts = (end_point_ts[:, :, :-1] +
                           end_point_ts[:, :, 1:]) / 2
        raypos = campos[:, None,
                        None, :] + raydir[:, :,
                                          None, :] * middle_point_ts[:, :, :,
                                                                     None]
        valid = torch.ones_like(middle_point_ts,
                                dtype=middle_point_ts.dtype,
                                device=middle_point_ts.device)

    return raypos, segment_length, valid, middle_point_ts
//  [python]
def test_normals_poly3D():
    vertices, faces = _create_simple_polygon()
    poly3D = TrimeshPolyhedron(R3('x'), vertices=vertices, faces=faces)    
    points = poly3D.boundary.sample_grid(n=15)
    normals = poly3D.boundary.normal(points)
    assert normals.shape == (15,3)
    for i in range(15):
        assert torch.isclose(normals[i].norm(), torch.tensor(1.0))
//  [python]
def test_normals_poly3D():
    vertices, faces = _create_simple_polygon()
    poly3D = TrimeshPolyhedron(R3('x'), vertices=vertices, faces=faces)    
    points = poly3D.boundary.sample_grid(n=15)
    normals = poly3D.boundary.normal(points)
    assert normals.shape == (15,3)
    for i in range(15):
        assert torch.isclose(torch.linalg.norm(normals[i]), torch.tensor(1.0))
//  [python]
def calc_procrustes_transform(
    X: Tensor, Y: Tensor
) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
    """Calculates procrustes transform of point clouds in batch format.
    minimize ||scale x  rot_mat x Y +t -X||_F with scale, rot_mat and translation
    code adapted from : http://stackoverflow.com/a/18927641/1884420
    Args:
        X (Tensor): batch x n x p
        Y (Tensor): batch x n x k
        Note: For joints n =21 and k=p=3

    Returns:
        y_transform (Tensor): transformed Y to best match X
        rot_mat (Tensor): Rotation matrix
        scale (Tensor): Scale
        translation (Tensor): Translation
    """
    if torch.all(X == 0):
        print("X contains only NaNs. Not computing PMSE.")
        return Y, (torch.tensor([]),) * 3
    if torch.all(Y == 0):
        print("Y contains only NaNs. Not computing PMSE.")
        return Y, (torch.tensor([]),) * 3
    with torch.no_grad():
        muX = X.mean(dim=1, keepdim=True)
        muY = Y.mean(dim=1, keepdim=True)
        # Centering and scale normalizing.
        X0 = X - muX
        Y0 = Y - muY
        normX = torch.norm(X0, dim=[1, 2], p="fro", keepdim=True)
        normY = torch.norm(Y0, dim=[1, 2], p="fro", keepdim=True)
        # Scale to equal (unit) norm
        X0 = X0 / normX
        Y0 = Y0 / normY
        # Compute optimum rotation matrix of Y
        A = torch.bmm(X0.transpose(2, 1), Y0)
        U, s, V = torch.svd(A)
        rot_mat = torch.bmm(V, U.transpose(2, 1))
        # Make sure we have a rotation
        det_rot_mat = torch.det(rot_mat)
        V[:, :, -1] *= torch.sign(det_rot_mat).view(-1, 1)
        s[:, -1] *= torch.sign(det_rot_mat)
        rot_mat = torch.matmul(V, U.transpose(2, 1))
        scale_ratio = s.sum(dim=1).view(-1, 1, 1)
        scale = scale_ratio * normX / normY
        translation = muX - scale * torch.matmul(muY, rot_mat)
        y_transform = normX * scale_ratio * torch.matmul(Y0, rot_mat) + muX
    return y_transform, rot_mat, scale, translation
//  [python]
def calc_procrustes_transform(
    X: Tensor, Y: Tensor
) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
    """Calculates procrustes transform of point clouds in batch format.
    minimize ||scale x  rot_mat x Y +t -X||_F with scale, rot_mat and translation
    code adapted from : http://stackoverflow.com/a/18927641/1884420
    Args:
        X (Tensor): batch x n x p
        Y (Tensor): batch x n x k
        Note: For joints n =21 and k=p=3

    Returns:
        y_transform (Tensor): transformed Y to best match X
        rot_mat (Tensor): Rotation matrix
        scale (Tensor): Scale
        translation (Tensor): Translation
    """
    if torch.all(X == 0):
        print("X contains only NaNs. Not computing PMSE.")
        return Y, (torch.tensor([]),) * 3
    if torch.all(Y == 0):
        print("Y contains only NaNs. Not computing PMSE.")
        return Y, (torch.tensor([]),) * 3
    with torch.no_grad():
        muX = X.mean(dim=1, keepdim=True)
        muY = Y.mean(dim=1, keepdim=True)
        # Centering and scale normalizing.
        X0 = X - muX
        Y0 = Y - muY
        normX = torch.linalg.norm(X0, dim=[1, 2], ord="fro", keepdim=True)
        normY = torch.linalg.norm(Y0, dim=[1, 2], ord="fro", keepdim=True)
        # Scale to equal (unit) norm
        X0 = X0 / normX
        Y0 = Y0 / normY
        # Compute optimum rotation matrix of Y
        A = torch.bmm(X0.transpose(2, 1), Y0)
        U, s, V = torch.svd(A)
        rot_mat = torch.bmm(V, U.transpose(2, 1))
        # Make sure we have a rotation
        det_rot_mat = torch.det(rot_mat)
        V[:, :, -1] *= torch.sign(det_rot_mat).view(-1, 1)
        s[:, -1] *= torch.sign(det_rot_mat)
        rot_mat = torch.matmul(V, U.transpose(2, 1))
        scale_ratio = s.sum(dim=1).view(-1, 1, 1)
        scale = scale_ratio * normX / normY
        translation = muX - scale * torch.matmul(muY, rot_mat)
        y_transform = normX * scale_ratio * torch.matmul(Y0, rot_mat) + muX
    return y_transform, rot_mat, scale, translation
//  [python]
def forward(self, positions):
        def contract(x):
            mag = torch.norm(x, p=self.order, dim=-1)[..., None]
            return torch.where(mag < 1, x, (2 - (1 / mag)) * (x / mag))

        if isinstance(positions, Gaussians):
            means = contract(positions.mean.clone())

            def contract_gauss(x):
                return (2 - 1 / torch.norm(x, p=self.order, dim=-1, keepdim=True)) * (
                    x / torch.norm(x, p=self.order, dim=-1, keepdim=True)
                )

            jc_means = vmap(jacrev(contract_gauss))(positions.mean.view(-1, positions.mean.shape[-1]))
            jc_means = jc_means.view(list(positions.mean.shape) + [positions.mean.shape[-1]])

            # Only update covariances on positions outside the unit sphere
            mag = positions.mean.norm(dim=-1)
            mask = mag >= 1
            cov = positions.cov.clone()
            cov[mask] = jc_means[mask] @ positions.cov[mask] @ torch.transpose(jc_means[mask], -2, -1)

            return Gaussians(mean=means, cov=cov)

        return contract(positions)
//  [python]
def forward(self, positions):
        def contract(x):
            mag = torch.linalg.norm(x, ord=self.order, dim=-1)[..., None]
            return torch.where(mag < 1, x, (2 - (1 / mag)) * (x / mag))

        if isinstance(positions, Gaussians):
            means = contract(positions.mean.clone())

            def contract_gauss(x):
                return (2 - 1 / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True)) * (
                    x / torch.linalg.norm(x, ord=self.order, dim=-1, keepdim=True)
                )

            jc_means = vmap(jacrev(contract_gauss))(positions.mean.view(-1, positions.mean.shape[-1]))
            jc_means = jc_means.view(list(positions.mean.shape) + [positions.mean.shape[-1]])

            # Only update covariances on positions outside the unit sphere
            mag = positions.mean.norm(dim=-1)
            mask = mag >= 1
            cov = positions.cov.clone()
            cov[mask] = jc_means[mask] @ positions.cov[mask] @ torch.transpose(jc_means[mask], -2, -1)

            return Gaussians(mean=means, cov=cov)

        return contract(positions)
//  [python]
def threshold_edges(self, x, threshold):
    # get mean attention
    # i) sparsify on S_hat
    if self.opt['new_edges'] == 'k_hop_att' and self.opt['sparsify'] == 'S_hat':
      attention_weights = self.odefunc.attention_weights
      mean_att = attention_weights
    # ii) sparsify on recalculated attentions
    else:#elif self.opt['sparsify'] == 'recalc_att':
      attention_weights = self.get_attention_weights(x)
      mean_att = attention_weights.mean(dim=1, keepdim=False)

    if self.opt['use_flux']:
      src_features = x[self.data_edge_index[0, :], :]
      dst_features = x[self.data_edge_index[1, :], :]
      delta = torch.norm(src_features - dst_features, dim=1)
      mean_att = mean_att * delta

    # just for the test where threshold catches all edges
    # unique_att = torch.unique(mean_att, sorted=False, return_inverse=False, return_counts=False, dim=0)
    # print(f"mean_att {mean_att.shape}, unqiue atts: {unique_att.shape}")
    # threshold

    # threshold = torch.quantile(mean_att, 1 - self.opt['att_samp_pct'])
    mask = mean_att > threshold
    self.odefunc.edge_index = self.data_edge_index[:, mask.T]
    sampled_attention_weights = self.renormalise_attention(mean_att[mask])
    print('retaining {} of {} edges'.format(self.odefunc.edge_index.shape[1], self.data_edge_index.shape[1]))
    self.data_edge_index = self.data_edge_index[:, mask.T]
    self.odefunc.edge_weight = sampled_attention_weights #rewiring structure so need to replace any preproc ew's with new ew's
    self.odefunc.attention_weights = sampled_attention_weights
//  [python]
def threshold_edges(self, x, threshold):
    # get mean attention
    # i) sparsify on S_hat
    if self.opt['new_edges'] == 'k_hop_att' and self.opt['sparsify'] == 'S_hat':
      attention_weights = self.odefunc.attention_weights
      mean_att = attention_weights
    # ii) sparsify on recalculated attentions
    else:#elif self.opt['sparsify'] == 'recalc_att':
      attention_weights = self.get_attention_weights(x)
      mean_att = attention_weights.mean(dim=1, keepdim=False)

    if self.opt['use_flux']:
      src_features = x[self.data_edge_index[0, :], :]
      dst_features = x[self.data_edge_index[1, :], :]
      delta = torch.linalg.norm(src_features - dst_features, dim=1)
      mean_att = mean_att * delta

    # just for the test where threshold catches all edges
    # unique_att = torch.unique(mean_att, sorted=False, return_inverse=False, return_counts=False, dim=0)
    # print(f"mean_att {mean_att.shape}, unqiue atts: {unique_att.shape}")
    # threshold

    # threshold = torch.quantile(mean_att, 1 - self.opt['att_samp_pct'])
    mask = mean_att > threshold
    self.odefunc.edge_index = self.data_edge_index[:, mask.T]
    sampled_attention_weights = self.renormalise_attention(mean_att[mask])
    print('retaining {} of {} edges'.format(self.odefunc.edge_index.shape[1], self.data_edge_index.shape[1]))
    self.data_edge_index = self.data_edge_index[:, mask.T]
    self.odefunc.edge_weight = sampled_attention_weights #rewiring structure so need to replace any preproc ew's with new ew's
    self.odefunc.attention_weights = sampled_attention_weights
//  [python]
def compute_random_projections(dim, device, n_random_projections=256):
    projs = torch.randn((n_random_projections, dim), device=device)
    projs /= torch.norm(projs, dim=-1, keepdim=True)
    return projs
//  [python]
def compute_random_projections(dim, device, n_random_projections=256):
    projs = torch.randn((n_random_projections, dim), device=device)
    projs /= torch.linalg.norm(projs, dim=-1)[...,None]
    return projs
//  [python]
def test_get_tile_bin_edges():
    from gsplat import _torch_impl
    from gsplat import get_tile_bin_edges

    torch.manual_seed(42)

    num_points = 100

    means3d = torch.randn((num_points, 3), device=device, requires_grad=True)
    scales = torch.randn((num_points, 3), device=device)
    glob_scale = 0.3
    quats = torch.randn((num_points, 4), device=device)
    quats /= torch.norm(quats, dim=-1, keepdim=True)
    viewmat = torch.eye(4, device=device)
    fx, fy = 3.0, 3.0
    H, W = 512, 512
    clip_thresh = 0.01

    BLOCK_SIZE = 16
    tile_bounds = (
        (W + BLOCK_SIZE - 1) // BLOCK_SIZE,
        (H + BLOCK_SIZE - 1) // BLOCK_SIZE,
        1,
    )

    (
        _cov3d,
        _cov2d,
        _xys,
        _depths,
        _radii,
        _conics,
        _compensation,
        _num_tiles_hit,
        _masks,
    ) = _torch_impl.project_gaussians_forward(
        means3d,
        scales,
        glob_scale,
        quats,
        viewmat,
        (fx, fy, W / 2, H / 2),
        (H, W),
        BLOCK_SIZE,
        clip_thresh,
    )

    _xys = _xys[_masks]
    _depths = _depths[_masks]
    _radii = _radii[_masks]
    _conics = _conics[_masks]
    _num_tiles_hit = _num_tiles_hit[_masks]

    num_points = num_points - torch.count_nonzero(~_masks).item()

    _cum_tiles_hit = torch.cumsum(_num_tiles_hit, dim=0, dtype=torch.int32)
    _num_intersects = _cum_tiles_hit[-1].item()
    _depths = _depths.contiguous()

    (
        _isect_ids_unsorted,
        _gaussian_ids_unsorted,
    ) = _torch_impl.map_gaussian_to_intersects(
        num_points, _xys, _depths, _radii, _cum_tiles_hit, tile_bounds, BLOCK_SIZE
    )

    # Sorting isect_ids_unsorted
    sorted_values, sorted_indices = torch.sort(_isect_ids_unsorted)

    _isect_ids_sorted = sorted_values
    _gaussian_ids_sorted = torch.gather(_gaussian_ids_unsorted, 0, sorted_indices)

    _tile_bins = _torch_impl.get_tile_bin_edges(
        _num_intersects, _isect_ids_sorted, tile_bounds
    )
    tile_bins = get_tile_bin_edges(_num_intersects, _isect_ids_sorted, tile_bounds)

    torch.testing.assert_close(_tile_bins, tile_bins)
//  [python]
def test_get_tile_bin_edges():
    from gsplat import _torch_impl
    from gsplat import get_tile_bin_edges

    torch.manual_seed(42)

    num_points = 100

    means3d = torch.randn((num_points, 3), device=device, requires_grad=True)
    scales = torch.randn((num_points, 3), device=device)
    glob_scale = 0.3
    quats = torch.randn((num_points, 4), device=device)
    quats /= torch.linalg.norm(quats, dim=-1, keepdim=True)
    viewmat = torch.eye(4, device=device)
    fx, fy = 3.0, 3.0
    H, W = 512, 512
    clip_thresh = 0.01

    BLOCK_SIZE = 16
    tile_bounds = (
        (W + BLOCK_SIZE - 1) // BLOCK_SIZE,
        (H + BLOCK_SIZE - 1) // BLOCK_SIZE,
        1,
    )

    (
        _cov3d,
        _cov2d,
        _xys,
        _depths,
        _radii,
        _conics,
        _compensation,
        _num_tiles_hit,
        _masks,
    ) = _torch_impl.project_gaussians_forward(
        means3d,
        scales,
        glob_scale,
        quats,
        viewmat,
        (fx, fy, W / 2, H / 2),
        (H, W),
        BLOCK_SIZE,
        clip_thresh,
    )

    _xys = _xys[_masks]
    _depths = _depths[_masks]
    _radii = _radii[_masks]
    _conics = _conics[_masks]
    _num_tiles_hit = _num_tiles_hit[_masks]

    num_points = num_points - torch.count_nonzero(~_masks).item()

    _cum_tiles_hit = torch.cumsum(_num_tiles_hit, dim=0, dtype=torch.int32)
    _num_intersects = _cum_tiles_hit[-1].item()
    _depths = _depths.contiguous()

    (
        _isect_ids_unsorted,
        _gaussian_ids_unsorted,
    ) = _torch_impl.map_gaussian_to_intersects(
        num_points, _xys, _depths, _radii, _cum_tiles_hit, tile_bounds, BLOCK_SIZE
    )

    # Sorting isect_ids_unsorted
    sorted_values, sorted_indices = torch.sort(_isect_ids_unsorted)

    _isect_ids_sorted = sorted_values
    _gaussian_ids_sorted = torch.gather(_gaussian_ids_unsorted, 0, sorted_indices)

    _tile_bins = _torch_impl.get_tile_bin_edges(
        _num_intersects, _isect_ids_sorted, tile_bounds
    )
    tile_bins = get_tile_bin_edges(_num_intersects, _isect_ids_sorted, tile_bounds)

    torch.testing.assert_close(_tile_bins, tile_bins)
//  [python]
def forward(self, x: torch.Tensor, y: torch.Tensor):
        """
        Arguments:
            x: features from one domain
            y: features from the other domain
        """
        embedding_size = x.shape[1]
        cx = covariance(x)
        cy = covariance(y)
        squared_fro_norm = torch.norm(cx - cy, p="fro") ** 2
        return squared_fro_norm / (4 * (embedding_size**2))
//  [python]
def forward(self, x: torch.Tensor, y: torch.Tensor):
        """
        Arguments:
            x: features from one domain
            y: features from the other domain
        """
        embedding_size = x.shape[1]
        cx = covariance(x)
        cy = covariance(y)
        squared_fro_norm = torch.linalg.norm(cx - cy, ord="fro") ** 2
        return squared_fro_norm / (4 * (embedding_size**2))
//  [python]
def get_local_tris(vertices, faces, device=torch.device("cpu")):
    fverts = vertices[faces].to(device)
    e1 = fverts[:, 1, :] - fverts[:, 0, :]
    e2 = fverts[:, 2, :] - fverts[:, 0, :]
    s = torch.norm(e1, dim=1)
    t = torch.norm(e2, dim=1)
    angle = torch.acos(torch.sum(e1 / s[:, None] * e2 / t[:, None], dim=1))
    x = torch.column_stack([torch.zeros(len(angle)).to(device), s, t * torch.cos(angle)])
    y = torch.column_stack([torch.zeros(len(angle)).to(device), torch.zeros(len(angle)).to(device), t * torch.sin(angle)])
    local_tris = torch.stack((x, y), dim=-1).reshape(len(angle), 3, 2)
    return local_tris
//  [python]
def get_local_tris(vertices, faces, device=torch.device("cpu")):
    fverts = vertices[faces].to(device)
    e1 = fverts[:, 1, :] - fverts[:, 0, :]
    e2 = fverts[:, 2, :] - fverts[:, 0, :]
    s = torch.linalg.norm(e1, dim=1)
    t = torch.linalg.norm(e2, dim=1)
    angle = torch.acos(torch.sum(e1 / s[:, None] * e2 / t[:, None], dim=1))
    x = torch.column_stack([torch.zeros(len(angle)).to(device), s, t * torch.cos(angle)])
    y = torch.column_stack([torch.zeros(len(angle)).to(device), torch.zeros(len(angle)).to(device), t * torch.sin(angle)])
    local_tris = torch.stack((x, y), dim=-1).reshape(len(angle), 3, 2)
    return local_tris
//  [python]
def map_txt_to_clip_feature(clip_model, tokenizer, prompt):
    text_inputs = tokenizer(
        prompt,
        padding="max_length",
        max_length=tokenizer.model_max_length,
        return_tensors="pt",
    )
    text_input_ids = text_inputs.input_ids

    if text_input_ids.shape[-1] > tokenizer.model_max_length:
        removed_text = tokenizer.batch_decode(text_input_ids[:, tokenizer.model_max_length :])
        logger.warning(
            "The following part of your input was truncated because CLIP can only handle sequences up to"
            f" {tokenizer.model_max_length} tokens: {removed_text}"
        )
        text_input_ids = text_input_ids[:, : tokenizer.model_max_length]
    text_embeddings = clip_model.get_text_features(text_input_ids.to(clip_model.device))
    text_embeddings = text_embeddings / torch.norm(text_embeddings, dim=-1, keepdim=True)
    text_embeddings = text_embeddings[:, None, :]
    return text_embeddings[0][0].cpu().detach().numpy()
//  [python]
def map_txt_to_clip_feature(clip_model, tokenizer, prompt):
    text_inputs = tokenizer(
        prompt,
        padding="max_length",
        max_length=tokenizer.model_max_length,
        return_tensors="pt",
    )
    text_input_ids = text_inputs.input_ids

    if text_input_ids.shape[-1] > tokenizer.model_max_length:
        removed_text = tokenizer.batch_decode(text_input_ids[:, tokenizer.model_max_length :])
        logger.warning(
            "The following part of your input was truncated because CLIP can only handle sequences up to"
            f" {tokenizer.model_max_length} tokens: {removed_text}"
        )
        text_input_ids = text_input_ids[:, : tokenizer.model_max_length]
    text_embeddings = clip_model.get_text_features(text_input_ids.to(clip_model.device))
    text_embeddings = text_embeddings / torch.linalg.norm(text_embeddings, dim=-1, keepdim=True)
    text_embeddings = text_embeddings[:, None, :]
    return text_embeddings[0][0].cpu().detach().numpy()
//  [python]
def map_img_to_model_feature(model, feature_extractor, imgs, device):
    for i, image in enumerate(imgs):
        if not image.mode == "RGB":
            imgs[i] = image.convert("RGB")
    imgs = normalize_images(imgs)
    retrieved_images = preprocess_images(imgs, feature_extractor).to(device)
    image_embeddings = model(retrieved_images)
    image_embeddings = image_embeddings / torch.norm(image_embeddings, dim=-1, keepdim=True)
    image_embeddings = image_embeddings[None, ...]
    return image_embeddings.cpu().detach().numpy()[0][0]
//  [python]
def map_img_to_model_feature(model, feature_extractor, imgs, device):
    for i, image in enumerate(imgs):
        if not image.mode == "RGB":
            imgs[i] = image.convert("RGB")
    imgs = normalize_images(imgs)
    retrieved_images = preprocess_images(imgs, feature_extractor).to(device)
    image_embeddings = model(retrieved_images)
    image_embeddings = image_embeddings / torch.linalg.norm(image_embeddings, dim=-1, keepdim=True)
    image_embeddings = image_embeddings[None, ...]
    return image_embeddings.cpu().detach().numpy()[0][0]
//  [python]
def solve(self, trust_radius):
        """Solve quadratic subproblem"""

        # Compute the Newton point.
        # This is the optimum for the quadratic model function.
        # If it is inside the trust radius then return this point.
        p_best = self.newton_point()
        if t.norm(p_best) < trust_radius:
            hits_boundary = False
            return p_best, hits_boundary

        # Compute the Cauchy point.
        # This is the predicted optimum along the direction of steepest descent.
        p_u = self.cauchy_point()

        # If the Cauchy point is outside the trust region,
        # then return the point where the path intersects the boundary.
        p_u_norm = t.norm(p_u)
        if p_u_norm >= trust_radius:
            p_boundary = p_u * (trust_radius / p_u_norm)
            hits_boundary = True
            return p_boundary, hits_boundary

        # Compute the intersection of the trust region boundary
        # and the line segment connecting the Cauchy and Newton points.
        # This requires solving a quadratic equation.
        # ||p_u + t*(p_best - p_u)||**2 == trust_radius**2
        # Solve this for positive time t using the quadratic formula.
        _, tb = self.get_boundaries_intersections(p_u, p_best - p_u,
                                                  trust_radius)
        p_boundary = p_u + tb * (p_best - p_u)
        hits_boundary = True
        return p_boundary, hits_boundary
//  [python]
def solve(self, trust_radius):
        """Solve quadratic subproblem"""

        # Compute the Newton point.
        # This is the optimum for the quadratic model function.
        # If it is inside the trust radius then return this point.
        p_best = self.newton_point()
        if norm(p_best) < trust_radius:
            hits_boundary = False
            return p_best, hits_boundary

        # Compute the Cauchy point.
        # This is the predicted optimum along the direction of steepest descent.
        p_u = self.cauchy_point()

        # If the Cauchy point is outside the trust region,
        # then return the point where the path intersects the boundary.
        p_u_norm = norm(p_u)
        if p_u_norm >= trust_radius:
            p_boundary = p_u * (trust_radius / p_u_norm)
            hits_boundary = True
            return p_boundary, hits_boundary

        # Compute the intersection of the trust region boundary
        # and the line segment connecting the Cauchy and Newton points.
        # This requires solving a quadratic equation.
        # ||p_u + t*(p_best - p_u)||**2 == trust_radius**2
        # Solve this for positive time t using the quadratic formula.
        _, tb = self.get_boundaries_intersections(p_u, p_best - p_u,
                                                  trust_radius)
        p_boundary = p_u + tb * (p_best - p_u)
        hits_boundary = True
        return p_boundary, hits_boundary
//  [python]
def update_symm_subs(xyz, pair, symmids, symmsub, symmRs, metasymm, fit=False, tscale=1.0):
    B,Ls = xyz.shape[0:2]
    Osub = symmsub.shape[0]
    L = Ls//Osub

    com = xyz[:,:L,1].mean(dim=-2)
    rcoms = torch.einsum('sij,bj->si', symmRs, com)
    subsymms, nneighs = metasymm
    symmsub_new = []
    for i in range(len(subsymms)):
        drcoms = torch.norm(rcoms[0,:] - rcoms[subsymms[i],:], dim=-1)
        _,subs_i = torch.topk(drcoms,nneighs[i],largest=False)
        subs_i,_ = torch.sort( subsymms[i][subs_i] )
        symmsub_new.append(subs_i)

    symmsub_new = torch.cat(symmsub_new)

    s_old = symmids[symmsub[:,None],symmsub[None,:]]
    s_new = symmids[symmsub_new[:,None],symmsub_new[None,:]]

    # remap old->new
    # a) find highest-magnitude patches
    pairsub = dict()
    pairmag = dict()
    for i in range(Osub):
        for j in range(Osub):
            idx_old = s_old[i,j].item()
            sub_ij = pair[:,i*L:(i+1)*L,j*L:(j+1)*L,:].clone()
            mag_ij = torch.max(sub_ij.flatten()) #torch.norm(sub_ij.flatten())
            if idx_old not in pairsub or mag_ij > pairmag[idx_old]:
                pairmag[idx_old] = mag_ij
                pairsub[idx_old] = (i,j) #sub_ij

    # b) reindex
    idx = torch.zeros((Osub*L,Osub*L),dtype=torch.long,device=pair.device)
    idx = (
        torch.arange(Osub*L,device=pair.device)[:,None]*Osub*L
         + torch.arange(Osub*L,device=pair.device)[None,:]
    )
    for i in range(Osub):
        for j in range(Osub):
            idx_new = s_new[i,j].item()
            if idx_new in pairsub:
                inew,jnew = pairsub[idx_new]
                idx[i*L:(i+1)*L,j*L:(j+1)*L] = (
                    Osub*L*torch.arange(inew*L,(inew+1)*L)[:,None]
                    + torch.arange(jnew*L,(jnew+1)*L)[None,:]
                )
    pair = pair.view(1,-1,pair.shape[-1])[:,idx.flatten(),:].view(1,Osub*L,Osub*L,pair.shape[-1])

    if symmsub is not None and symmsub.shape[0]>1:
        xyz = update_symm_Rs(xyz, L, symmsub_new, symmRs, fit, tscale)

    return xyz, pair, symmsub_new
//  [python]
def update_symm_subs(xyz, pair, symmids, symmsub, symmRs, metasymm, fit=False, tscale=1.0):
    B,Ls = xyz.shape[0:2]
    Osub = symmsub.shape[0]
    L = Ls//Osub

    com = xyz[:,:L,1].mean(dim=-2)
    rcoms = torch.einsum('sij,bj->si', symmRs, com)
    subsymms, nneighs = metasymm
    symmsub_new = []
    for i in range(len(subsymms)):
        drcoms = torch.linalg.norm(rcoms[0,:] - rcoms[subsymms[i],:], dim=-1)
        _,subs_i = torch.topk(drcoms,nneighs[i],largest=False)
        subs_i,_ = torch.sort( subsymms[i][subs_i] )
        symmsub_new.append(subs_i)

    symmsub_new = torch.cat(symmsub_new)

    s_old = symmids[symmsub[:,None],symmsub[None,:]]
    s_new = symmids[symmsub_new[:,None],symmsub_new[None,:]]

    # remap old->new
    # a) find highest-magnitude patches
    pairsub = dict()
    pairmag = dict()
    for i in range(Osub):
        for j in range(Osub):
            idx_old = s_old[i,j].item()
            sub_ij = pair[:,i*L:(i+1)*L,j*L:(j+1)*L,:].clone()
            mag_ij = torch.max(sub_ij.flatten()) #torch.norm(sub_ij.flatten())
            if idx_old not in pairsub or mag_ij > pairmag[idx_old]:
                pairmag[idx_old] = mag_ij
                pairsub[idx_old] = (i,j) #sub_ij

    # b) reindex
    idx = torch.zeros((Osub*L,Osub*L),dtype=torch.long,device=pair.device)
    idx = (
        torch.arange(Osub*L,device=pair.device)[:,None]*Osub*L
         + torch.arange(Osub*L,device=pair.device)[None,:]
    )
    for i in range(Osub):
        for j in range(Osub):
            idx_new = s_new[i,j].item()
            if idx_new in pairsub:
                inew,jnew = pairsub[idx_new]
                idx[i*L:(i+1)*L,j*L:(j+1)*L] = (
                    Osub*L*torch.arange(inew*L,(inew+1)*L)[:,None]
                    + torch.arange(jnew*L,(jnew+1)*L)[None,:]
                )
    pair = pair.view(1,-1,pair.shape[-1])[:,idx.flatten(),:].view(1,Osub*L,Osub*L,pair.shape[-1])

    if symmsub is not None and symmsub.shape[0]>1:
        xyz = update_symm_Rs(xyz, L, symmsub_new, symmRs, fit, tscale)

    return xyz, pair, symmsub_new
//  [python]
def update_symm_Rs(xyz, Lasu, symmsub, symmRs, fit=False, tscale=1.0):
    def dist_error_comp(R0,T0,xyz,tscale):
        B = xyz.shape[0]
        Tcom = xyz[:,:Lasu].mean(dim=1,keepdim=True)
        Tcorr = torch.einsum('ij,brj->bri', R0, xyz[:,:Lasu]-Tcom) + Tcom + tscale*T0[None,None,:]

        # distance map loss
        Xsymm = torch.einsum('sij,brj->bsri', symmRs[symmsub], Tcorr).reshape(B,-1,3)
        Xtrue = Ts

        delsx = Xsymm[:,:Lasu,None]-Xsymm[:, None, Lasu:]
        deltx = Xtrue[:,:Lasu,None]-Xtrue[:, None, Lasu:]
        dsymm = torch.norm(delsx, dim=-1)
        dtrue = torch.norm(deltx, dim=-1)
        loss1 = torch.abs(dsymm-dtrue).mean()

        return loss1,0.0 #loss2

    def dist_error(R0,T0,xyz,tscale,w_clash=0.0):
        l1,l2 = dist_error_comp(R0,T0,xyz,tscale)
        return l1+w_clash*l2

    def Q2R(Q):
        Qs = torch.cat((torch.ones((1),device=Q.device),Q),dim=-1)
        Qs = normQ(Qs)
        return Qs2Rs(Qs[None,:]).squeeze(0)
        

    B = xyz.shape[0]

    # symmetry correction 1: don't let COM (of entire complex) move
    #Tmean = xyz[:,:Lasu].reshape(-1,3).mean(dim=0)
    #Tmean = torch.einsum('sij,j->si', symmRs, Tmean).mean(dim=0)
    #xyz = xyz - Tmean

    if fit:
        # symmetry correction 2: use minimization to minimize drms
        with torch.enable_grad():
            T0 = torch.zeros(3,device=xyz.device).requires_grad_(True)
            Q0 = torch.zeros(3,device=xyz.device).requires_grad_(True)
            lbfgs = torch.optim.LBFGS([T0,Q0],
                        history_size=10, 
                        max_iter=4,
                        line_search_fn="strong_wolfe")
            def closure():
                lbfgs.zero_grad()
                loss = dist_error(Q2R(Q0),T0,xyz[:,:,1], tscale)
                loss.backward() #retain_graph=True)
                return loss

            for e in range(4):
                loss = lbfgs.step(closure)

            Tcom = xyz[:,:Lasu].mean(dim=1,keepdim=True).detach()
            Q0 = Q0.detach()
            T0 = T0.detach()
            xyz = torch.einsum('ij,braj->brai', Q2R(Q0), xyz[:,:Lasu]-Tcom) +Tcom + tscale*T0[None,None,:]

    xyz = torch.einsum('sij,braj->bsrai', symmRs[symmsub], xyz[:,:Lasu])
    xyz = xyz.reshape(B,-1,3,3) # (B,S,L,3,3)

    return xyz
//  [python]
def update_symm_Rs(xyz, Lasu, symmsub, symmRs, fit=False, tscale=1.0):
    def dist_error_comp(R0,T0,xyz,tscale):
        B = xyz.shape[0]
        Tcom = xyz[:,:Lasu].mean(dim=1,keepdim=True)
        Tcorr = torch.einsum('ij,brj->bri', R0, xyz[:,:Lasu]-Tcom) + Tcom + tscale*T0[None,None,:]

        # distance map loss
        Xsymm = torch.einsum('sij,brj->bsri', symmRs[symmsub], Tcorr).reshape(B,-1,3)
        Xtrue = Ts

        delsx = Xsymm[:,:Lasu,None]-Xsymm[:, None, Lasu:]
        deltx = Xtrue[:,:Lasu,None]-Xtrue[:, None, Lasu:]
        dsymm = torch.linalg.norm(delsx, dim=-1)
        dtrue = torch.linalg.norm(deltx, dim=-1)
        loss1 = torch.abs(dsymm-dtrue).mean()

        return loss1,0.0 #loss2

    def dist_error(R0,T0,xyz,tscale,w_clash=0.0):
        l1,l2 = dist_error_comp(R0,T0,xyz,tscale)
        return l1+w_clash*l2

    def Q2R(Q):
        Qs = torch.cat((torch.ones((1),device=Q.device),Q),dim=-1)
        Qs = normQ(Qs)
        return Qs2Rs(Qs[None,:]).squeeze(0)
        

    B = xyz.shape[0]

    # symmetry correction 1: don't let COM (of entire complex) move
    #Tmean = xyz[:,:Lasu].reshape(-1,3).mean(dim=0)
    #Tmean = torch.einsum('sij,j->si', symmRs, Tmean).mean(dim=0)
    #xyz = xyz - Tmean

    if fit:
        # symmetry correction 2: use minimization to minimize drms
        with torch.enable_grad():
            T0 = torch.zeros(3,device=xyz.device).requires_grad_(True)
            Q0 = torch.zeros(3,device=xyz.device).requires_grad_(True)
            lbfgs = torch.optim.LBFGS([T0,Q0],
                        history_size=10, 
                        max_iter=4,
                        line_search_fn="strong_wolfe")
            def closure():
                lbfgs.zero_grad()
                loss = dist_error(Q2R(Q0),T0,xyz[:,:,1], tscale)
                loss.backward() #retain_graph=True)
                return loss

            for e in range(4):
                loss = lbfgs.step(closure)

            Tcom = xyz[:,:Lasu].mean(dim=1,keepdim=True).detach()
            Q0 = Q0.detach()
            T0 = T0.detach()
            xyz = torch.einsum('ij,braj->brai', Q2R(Q0), xyz[:,:Lasu]-Tcom) +Tcom + tscale*T0[None,None,:]

    xyz = torch.einsum('sij,braj->bsrai', symmRs[symmsub], xyz[:,:Lasu])
    xyz = xyz.reshape(B,-1,3,3) # (B,S,L,3,3)

    return xyz
//  [python]
def compute_velocity_loss(inputs: Dict[Any, Tensor],
                          outputs: Dict[Any, Tensor],
                          distances: Dict[float, float],
                          device: torch.device,
                          ) -> Tensor:
    batch_size = inputs[0].shape[0]  # might be different from self.batch_size
    velocity_loss = torch.zeros(batch_size, device=device).squeeze()
    num_frames = 0
    for frame in [0, -1, 1]:
        if frame == -1:
            continue
        if frame == 0:
            pred_translation = outputs[('translation', 0, -1)]
        else:
            pred_translation = outputs[('translation', 0, 1)]
        gt_distance = torch.abs(distances[frame]).squeeze()
        pred_distance = torch.norm(pred_translation, dim=-1).squeeze()
        velocity_loss += F.l1_loss(pred_distance,
                                   gt_distance,
                                   reduction='none')  # separated by sample in batch
        num_frames += 1
    velocity_loss /= num_frames
    return velocity_loss
//  [python]
def compute_velocity_loss(inputs: Dict[Any, Tensor],
                          outputs: Dict[Any, Tensor],
                          distances: Dict[float, float],
                          device: torch.device,
                          ) -> Tensor:
    batch_size = inputs[0].shape[0]  # might be different from self.batch_size
    velocity_loss = torch.zeros(batch_size, device=device).squeeze()
    num_frames = 0
    for frame in [0, -1, 1]:
        if frame == -1:
            continue
        if frame == 0:
            pred_translation = outputs[('translation', 0, -1)]
        else:
            pred_translation = outputs[('translation', 0, 1)]
        gt_distance = torch.abs(distances[frame]).squeeze()
        pred_distance = torch.linalg.norm(pred_translation, dim=-1).squeeze()
        velocity_loss += F.l1_loss(pred_distance,
                                   gt_distance,
                                   reduction='none')  # separated by sample in batch
        num_frames += 1
    velocity_loss /= num_frames
    return velocity_loss
//  [python]
def compute_reward(self, state: Tensor, action: Tensor) -> float:
    agent_atom = state if self.state_only else torch.cat([state, action], dim=1)
    agent_atom = self.data_scale * (agent_atom + self.data_offset)  # Normalise the agent atom
    weight, cost = 1 / self.time_horizon - 1e-6, 0  # Note: subtracting eps from initial agent atom weight for numerical stability
    dists = torch.norm(self.expert_atoms - agent_atom, dim=1)
    while weight > 0:
      closest_expert_idx = dists.argmin().item()  # Find closest expert atom
      expert_weight = self.expert_weights[closest_expert_idx].item()
      # Update costs and weights
      if weight >= expert_weight:
        cost += expert_weight * dists[closest_expert_idx].item()
        weight -= expert_weight
        self.expert_atoms, self.expert_weights, dists = _delete_row(self.expert_atoms, closest_expert_idx), _delete_row(self.expert_weights, closest_expert_idx), _delete_row(dists, closest_expert_idx)  # Remove the expert atom
      else:
        cost += weight * dists[closest_expert_idx].item()
        self.expert_weights[closest_expert_idx] -= weight
        weight = 0
    return self.reward_scale * exp(-self.reward_bandwidth * cost)