jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_model_un_normalization(self):
        names = self.open_clip_test_model

        sentences = ['hello', 'this is a test sentence. so is this.', ['hello', 'this is a test sentence. so is this.']]
        device = 'cpu'
        eps = 1e-3

        for name in names:
            # this model always give normalized output
            if name in ['open_clip/coca_ViT-B-32/mscoco_finetuned_laion2b_s13b_b90k',]:
                continue
            model_properties = get_model_properties_from_registry(name)
            model = _load_model(model_properties['name'], model_properties=model_properties, device=device)

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

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

            clear_loaded_models()
//  [python]
def test_model_un_normalization(self):
        names = self.open_clip_test_model

        sentences = ['hello', 'this is a test sentence. so is this.', ['hello', 'this is a test sentence. so is this.']]
        device = 'cpu'
        eps = 1e-3

        for name in names:
            # this model always give normalized output
            if name in ['open_clip/coca_ViT-B-32/mscoco_finetuned_laion2b_s13b_b90k',]:
                continue
            model_properties = get_model_properties_from_registry(name)
            model = _load_model(model_properties['name'], model_properties=model_properties, device=device)

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

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

            clear_loaded_models()
//  [python]
def test_model_un_normalization(self):
        # note: sbert native seems to provide normalized embeddings even with = False, needs more investigation
        # , 
        names = [ 'RN50', "ViT-B/16", "hf/all-MiniLM-L6-v1",
                 "hf/all_datasets_v4_MiniLM-L6", "onnx/all-MiniLM-L6-v1", "onnx/all_datasets_v4_MiniLM-L6"]

        names_e5 = ["hf/e5-small", "hf/e5-base", "hf/e5-small-unsupervised", "hf/e5-base-unsupervised", "hf/e5-base-v2",
                    "hf/multilingual-e5-small"]
        names += names_e5

        sentences = ['hello', 'this is a test sentence. so is this.', ['hello', 'this is a test sentence. so is this.']]
        device = 'cpu'
        eps = 1e-3

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

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

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

            clear_loaded_models()
//  [python]
def test_model_un_normalization(self):
        # note: sbert native seems to provide normalized embeddings even with = False, needs more investigation
        # , 
        names = [ 'RN50', "ViT-B/16", "hf/all-MiniLM-L6-v1",
                 "hf/all_datasets_v4_MiniLM-L6", "onnx/all-MiniLM-L6-v1", "onnx/all_datasets_v4_MiniLM-L6"]

        names_e5 = ["hf/e5-small", "hf/e5-base", "hf/e5-small-unsupervised", "hf/e5-base-unsupervised", "hf/e5-base-v2",
                    "hf/multilingual-e5-small"]
        names += names_e5

        sentences = ['hello', 'this is a test sentence. so is this.', ['hello', 'this is a test sentence. so is this.']]
        device = 'cpu'
        eps = 1e-3

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

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

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

            clear_loaded_models()
//  [python]
def test_k2_speech_recognition_on_the_fly_feature_extraction(
    k2_cut_set, use_batch_extract, fault_tolerant
):
    precomputed_dataset = K2SpeechRecognitionDataset()
    on_the_fly_dataset = K2SpeechRecognitionDataset(
        input_strategy=OnTheFlyFeatures(
            Fbank(FbankConfig(num_mel_bins=40)),
            use_batch_extract=use_batch_extract,
            fault_tolerant=fault_tolerant,
        )
    )
    sampler = SimpleCutSampler(k2_cut_set, shuffle=False, max_cuts=1)
    for cut_ids in sampler:
        batch_pc = precomputed_dataset[cut_ids]
        batch_otf = on_the_fly_dataset[cut_ids]

        # Check that the features do not differ too much.
        norm_pc = torch.norm(batch_pc["inputs"])
        norm_diff = torch.norm(batch_pc["inputs"] - batch_otf["inputs"])
        # The precomputed and on-the-fly features are different due to mixing in time/fbank domains
        # and lilcom compression.
        assert norm_diff < 0.01 * norm_pc

        # Check that the supervision boundaries are the same.
        assert (
            batch_pc["supervisions"]["start_frame"]
            == batch_otf["supervisions"]["start_frame"]
        ).all()
        assert (
            batch_pc["supervisions"]["num_frames"]
            == batch_otf["supervisions"]["num_frames"]
        ).all()
//  [python]
def test_k2_speech_recognition_on_the_fly_feature_extraction(
    k2_cut_set, use_batch_extract, fault_tolerant
):
    precomputed_dataset = K2SpeechRecognitionDataset()
    on_the_fly_dataset = K2SpeechRecognitionDataset(
        input_strategy=OnTheFlyFeatures(
            Fbank(FbankConfig(num_mel_bins=40)),
            use_batch_extract=use_batch_extract,
            fault_tolerant=fault_tolerant,
        )
    )
    sampler = SimpleCutSampler(k2_cut_set, shuffle=False, max_cuts=1)
    for cut_ids in sampler:
        batch_pc = precomputed_dataset[cut_ids]
        batch_otf = on_the_fly_dataset[cut_ids]

        # Check that the features do not differ too much.
        norm_pc = torch.linalg.norm(batch_pc["inputs"])
        norm_diff = torch.linalg.norm(batch_pc["inputs"] - batch_otf["inputs"])
        # The precomputed and on-the-fly features are different due to mixing in time/fbank domains
        # and lilcom compression.
        assert norm_diff < 0.01 * norm_pc

        # Check that the supervision boundaries are the same.
        assert (
            batch_pc["supervisions"]["start_frame"]
            == batch_otf["supervisions"]["start_frame"]
        ).all()
        assert (
            batch_pc["supervisions"]["num_frames"]
            == batch_otf["supervisions"]["num_frames"]
        ).all()
//  [python]
def get_polyline_length(polyline):
    """Calculate the length of a polyline.
    Args:
        polyline: Numpy array of shape (N,3)
    Returns:
        The length of the polyline as a scalar.
    Raises:
        RuntimeError: If `polyline` doesn't have shape (N,2) or (N,3).
    """
    if polyline.shape[1] not in [2, 3]:
        raise RuntimeError("Polyline must have shape (N,2) or (N,3)")
    offsets = torch.diff(polyline, axis=0)
    return float(torch.norm(offsets, dim=1).sum())
//  [python]
def get_polyline_length(polyline):
    """Calculate the length of a polyline.
    Args:
        polyline: Numpy array of shape (N,3)
    Returns:
        The length of the polyline as a scalar.
    Raises:
        RuntimeError: If `polyline` doesn't have shape (N,2) or (N,3).
    """
    if polyline.shape[1] not in [2, 3]:
        raise RuntimeError("Polyline must have shape (N,2) or (N,3)")
    offsets = torch.diff(polyline, axis=0)
    return float(torch.linalg.norm(offsets, axis=1).sum())
//  [python]
def get_norm(self, weight):
        with torch.no_grad():
            weight_q, weight_k = weight[:self.num_features, :], weight[self.num_features:, :]
            dot_matrix = weight_q @ weight_k.T * self.scale
            self.lipnorm = torch.norm(torch.eye(self.num_features,
                                                       device=weight.device) + 2 * dot_matrix,
                                             ord=2) + 5
//  [python]
def get_norm(self, weight):
        with torch.no_grad():
            weight_q, weight_k = weight[:self.num_features, :], weight[self.num_features:, :]
            dot_matrix = weight_q @ weight_k.T * self.scale
            self.lipnorm = torch.linalg.norm(torch.eye(self.num_features,
                                                       device=weight.device) + 2 * dot_matrix,
                                             ord=2) + 5
//  [python]
def evaluate_registration(self, output_dict, data_dict):
        transform = data_dict['transform']
        est_transform = output_dict['estimated_transform']
        src_points = output_dict['src_points']

        rre, rte = isotropic_transform_error(transform, est_transform)
        recall = torch.logical_and(torch.lt(rre, self.acceptance_rre), torch.lt(rte, self.acceptance_rte)).float()

        gt_src_points = apply_transform(src_points, transform)
        est_src_points = apply_transform(src_points, est_transform)
        rmse = torch.norm(est_src_points - gt_src_points, dim=1).mean()

        return rre, rte, rmse, recall
//  [python]
def evaluate_registration(self, output_dict, data_dict):
        transform = data_dict['transform']
        est_transform = output_dict['estimated_transform']
        src_points = output_dict['src_points']

        rre, rte = isotropic_transform_error(transform, est_transform)
        recall = torch.logical_and(torch.lt(rre, self.acceptance_rre), torch.lt(rte, self.acceptance_rte)).float()

        gt_src_points = apply_transform(src_points, transform)
        est_src_points = apply_transform(src_points, est_transform)
        rmse = torch.linalg.norm(est_src_points - gt_src_points, dim=1).mean()

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

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

        # corner case: no correct coarse match found
        if correct_mask.sum() == 0:
            if (
                self.training
            ):  # this seldomly happen when training, since we pad prediction with gt
                logger.warning("assign a false supervision to avoid ddp deadlock")
                correct_mask[0] = True
                weight[0] = 1e-6
            else:
                return None

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

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

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

        # corner case: no correct coarse match found
        if correct_mask.sum() == 0:
            if (
                self.training
            ):  # this seldomly happen when training, since we pad prediction with gt
                logger.warning("assign a false supervision to avoid ddp deadlock")
                correct_mask[0] = True
                weight[0] = 1e-6
            else:
                return None

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

        return loss
//  [python]
def distance_with(self, tgt_pos):
        """Compute the distance with an other pos

        Parameters
        ----------
        tgt_pos: aloscene.Pose
            Target pose to compute the distance with

        Returns
        -------
        n_pos: torch.tensor
            Translation tensor of shape (..., 3)
        """
        return torch.norm(self.translation_with(tgt_pos), dim=-1)
//  [python]
def distance_with(self, tgt_pos):
        """Compute the distance with an other pos

        Parameters
        ----------
        tgt_pos: aloscene.Pose
            Target pose to compute the distance with

        Returns
        -------
        n_pos: torch.tensor
            Translation tensor of shape (..., 3)
        """
        return torch.linalg.norm(self.translation_with(tgt_pos), dim=-1)
//  [python]
def iterative_pinv_analysis(
    identity_tolerance: float = 1e-1,
    pinv_tolerance: float = 5e-1,
    max_iters: int = 30,
    plot: bool = True,
):

    for i in range(1, 10):
        B, M = 1, 2**i
        a = torch.rand(B, M, M)
        a = torch.softmax(a, dim=-1)

        for n_iter in range(1, max_iters + 1):
            result = iterative_pinv(a, n_iter=n_iter)
            expected = torch.linalg.pinv(a)

            result_identity = torch.matmul(a, result)
            identity = torch.eye(M)

            # Default is frobenius norm.
            identity_error = torch.norm(identity - result_identity, dim=(-2, -1))
            inverse_error = torch.norm(expected - result, dim=(-2, -1))

            if (identity_error < identity_tolerance).all() or n_iter == max_iters:
                print(
                    f"Size {M}, n_iters {n_iter}: \n\t \
                    Final Error from Identity: {identity_error.item()} \n\t \
                    Final Error from linalg.pinv {inverse_error.item()}"
                )
                break
//  [python]
def iterative_pinv_analysis(
    identity_tolerance: float = 1e-1,
    pinv_tolerance: float = 5e-1,
    max_iters: int = 30,
    plot: bool = True,
):

    for i in range(1, 10):
        B, M = 1, 2**i
        a = torch.rand(B, M, M)
        a = torch.softmax(a, dim=-1)

        for n_iter in range(1, max_iters + 1):
            result = iterative_pinv(a, n_iter=n_iter)
            expected = torch.linalg.pinv(a)

            result_identity = torch.matmul(a, result)
            identity = torch.eye(M)

            # Default is frobenius norm.
            identity_error = torch.linalg.norm(identity - result_identity, dim=(-2, -1))
            inverse_error = torch.linalg.norm(expected - result, dim=(-2, -1))

            if (identity_error < identity_tolerance).all() or n_iter == max_iters:
                print(
                    f"Size {M}, n_iters {n_iter}: \n\t \
                    Final Error from Identity: {identity_error.item()} \n\t \
                    Final Error from linalg.pinv {inverse_error.item()}"
                )
                break
//  [python]
def nonlinear(self, Bz):
        norm = torch.norm(Bz.reshape(Bz.shape[0], Bz.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).unsqueeze(-1)
        pred = pred.unsqueeze(-1)
        out = torch.sum(gam * Bz * pred, axis=0)
        return out, y
//  [python]
def nonlinear(self, Bz):
        norm = torch.linalg.norm(Bz.reshape(Bz.shape[0], Bz.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).unsqueeze(-1)
        pred = pred.unsqueeze(-1)
        out = torch.sum(gam * Bz * pred, axis=0)
        return out, y
//  [python]
def generate_perspective_rays(camera: Camera, pixel_grid: Tuple[torch.Tensor, torch.Tensor]):
    # coords_grid should remain immutable (a new tensor is implicitly created here)
    pixel_y, pixel_x = pixel_grid
    pixel_x = pixel_x.to(camera.device, camera.dtype)
    pixel_y = pixel_y.to(camera.device, camera.dtype)

    # Account for principal point offset from canvas center
    pixel_x = pixel_x - camera.x0
    pixel_y = pixel_y + camera.y0

    # pixel values are now in range [-1, 1], both tensors are of shape res_y x res_x
    pixel_x = 2 * (pixel_x / camera.width) - 1.0
    pixel_y = 2 * (pixel_y / camera.height) - 1.0

    ray_dir = torch.stack((pixel_x * camera.tan_half_fov(CameraFOV.HORIZONTAL),
                           -pixel_y * camera.tan_half_fov(CameraFOV.VERTICAL),
                           -torch.ones_like(pixel_x)), dim=-1)

    ray_dir = ray_dir.reshape(-1, 3)    # Flatten grid rays to 1D array
    ray_orig = torch.zeros_like(ray_dir)

    # Transform from camera to world coordinates
    ray_orig, ray_dir = camera.extrinsics.inv_transform_rays(ray_orig, ray_dir)
    ray_dir /= torch.norm(ray_dir, dim=-1, keepdim=True)
    ray_orig, ray_dir = ray_orig[0], ray_dir[0]  # Assume a single camera

    return ray_orig, ray_dir, camera.near, camera.far
//  [python]
def generate_perspective_rays(camera: Camera, pixel_grid: Tuple[torch.Tensor, torch.Tensor]):
    # coords_grid should remain immutable (a new tensor is implicitly created here)
    pixel_y, pixel_x = pixel_grid
    pixel_x = pixel_x.to(camera.device, camera.dtype)
    pixel_y = pixel_y.to(camera.device, camera.dtype)

    # Account for principal point offset from canvas center
    pixel_x = pixel_x - camera.x0
    pixel_y = pixel_y + camera.y0

    # pixel values are now in range [-1, 1], both tensors are of shape res_y x res_x
    pixel_x = 2 * (pixel_x / camera.width) - 1.0
    pixel_y = 2 * (pixel_y / camera.height) - 1.0

    ray_dir = torch.stack((pixel_x * camera.tan_half_fov(CameraFOV.HORIZONTAL),
                           -pixel_y * camera.tan_half_fov(CameraFOV.VERTICAL),
                           -torch.ones_like(pixel_x)), dim=-1)

    ray_dir = ray_dir.reshape(-1, 3)    # Flatten grid rays to 1D array
    ray_orig = torch.zeros_like(ray_dir)

    # Transform from camera to world coordinates
    ray_orig, ray_dir = camera.extrinsics.inv_transform_rays(ray_orig, ray_dir)
    ray_dir /= torch.linalg.norm(ray_dir, dim=-1, keepdim=True)
    ray_orig, ray_dir = ray_orig[0], ray_dir[0]  # Assume a single camera

    return ray_orig, ray_dir, camera.near, camera.far
//  [python]
def py_given_x0(self, x0, y, A, verbose=False):
        norm_const = 1/((2 * np.pi)**self.input_dim * self.sigma_0**2)
        exp_in = -1/(2 * self.sigma_0**2) * torch.norm(y - A(x0))**2
        if not verbose:
            return norm_const * torch.exp(exp_in)
        else:
            return norm_const * torch.exp(exp_in), norm_const, exp_in
//  [python]
def py_given_x0(self, x0, y, A, verbose=False):
        norm_const = 1/((2 * np.pi)**self.input_dim * self.sigma_0**2)
        exp_in = -1/(2 * self.sigma_0**2) * torch.linalg.norm(y - A(x0))**2
        if not verbose:
            return norm_const * torch.exp(exp_in)
        else:
            return norm_const * torch.exp(exp_in), norm_const, exp_in
//  [python]
def pxt_given_x0(self, x0, xt, t, verbose=False):
        beta_t = self.betas[t]
        norm_const = 1/((2 * np.pi)**self.label_dim * beta_t)
        exp_in = -1/(2 * beta_t) * torch.norm(xt - np.sqrt(1 - beta_t)*x0)**2
        if not verbose:
            return norm_const * torch.exp(exp_in)
        else:
            return norm_const * torch.exp(exp_in), norm_const, exp_in
//  [python]
def pxt_given_x0(self, x0, xt, t, verbose=False):
        beta_t = self.betas[t]
        norm_const = 1/((2 * np.pi)**self.label_dim * beta_t)
        exp_in = -1/(2 * beta_t) * torch.linalg.norm(xt - np.sqrt(1 - beta_t)*x0)**2
        if not verbose:
            return norm_const * torch.exp(exp_in)
        else:
            return norm_const * torch.exp(exp_in), norm_const, exp_in
//  [python]
def volumetric_rendering(rgb, density, t_samples, dirs, white_bkgd):
    """Volumetric Rendering Function.
    Args:
        rgb: torch.Tensor, color, [batch_size, num_samples, 3]
        density: torch.Tensor, density, [batch_size, num_samples, 1].
        t_samples: torch.Tensor, [batch_size, num_samples+1].
        dirs: torch.Tensor, [batch_size, 3].
        white_bkgd: bool.
    Returns:
        comp_rgb: torch.Tensor, [batch_size, 3].
        disp: torch.Tensor, [batch_size].
        acc: torch.Tensor, [batch_size].
        weights: torch.Tensor, [batch_size, num_samples]
    """
    t_mids = 0.5 * (t_samples[..., :-1] + t_samples[..., 1:])
    t_interval = t_samples[..., 1:] - t_samples[..., :-1]  # [B, N]
    # models/mip.py:8 here sample point by multiply the interval with the direction without normalized, so
    # the delta is norm(t1*d-t2*d) = (t1-t2)*norm(d)
    delta = t_interval * torch.norm(torch.unsqueeze(dirs, dim=-2), dim=-1)
    # Note that we're quietly turning density from [..., 0] to [...].
    density_delta = density[..., 0] * delta

    alpha = 1 - torch.exp(-density_delta)
    trans = torch.exp(-torch.cat([
        torch.zeros_like(density_delta[..., :1]),
        torch.cumsum(density_delta[..., :-1], dim=-1)
    ], dim=-1))
    weights = alpha * trans  # [B, N]

    comp_rgb = (torch.unsqueeze(weights, dim=-1) * rgb).sum(axis=-2)  # [B, N, 1] * [B, N, 3] -> [B, 3]
    acc = weights.sum(axis=-1)
    distance = (weights * t_mids).sum(axis=-1)
    distance = torch.clamp(torch.nan_to_num(distance), t_samples[:, 0], t_samples[:, -1])
    if white_bkgd:
        comp_rgb = comp_rgb + (1. - torch.unsqueeze(acc, dim=-1))
    return comp_rgb, distance, acc, weights
//  [python]
def volumetric_rendering(rgb, density, t_samples, dirs, white_bkgd):
    """Volumetric Rendering Function.
    Args:
        rgb: torch.Tensor, color, [batch_size, num_samples, 3]
        density: torch.Tensor, density, [batch_size, num_samples, 1].
        t_samples: torch.Tensor, [batch_size, num_samples+1].
        dirs: torch.Tensor, [batch_size, 3].
        white_bkgd: bool.
    Returns:
        comp_rgb: torch.Tensor, [batch_size, 3].
        disp: torch.Tensor, [batch_size].
        acc: torch.Tensor, [batch_size].
        weights: torch.Tensor, [batch_size, num_samples]
    """
    t_mids = 0.5 * (t_samples[..., :-1] + t_samples[..., 1:])
    t_interval = t_samples[..., 1:] - t_samples[..., :-1]  # [B, N]
    # models/mip.py:8 here sample point by multiply the interval with the direction without normalized, so
    # the delta is norm(t1*d-t2*d) = (t1-t2)*norm(d)
    delta = t_interval * torch.linalg.norm(torch.unsqueeze(dirs, dim=-2), dim=-1)
    # Note that we're quietly turning density from [..., 0] to [...].
    density_delta = density[..., 0] * delta

    alpha = 1 - torch.exp(-density_delta)
    trans = torch.exp(-torch.cat([
        torch.zeros_like(density_delta[..., :1]),
        torch.cumsum(density_delta[..., :-1], dim=-1)
    ], dim=-1))
    weights = alpha * trans  # [B, N]

    comp_rgb = (torch.unsqueeze(weights, dim=-1) * rgb).sum(axis=-2)  # [B, N, 1] * [B, N, 3] -> [B, 3]
    acc = weights.sum(axis=-1)
    distance = (weights * t_mids).sum(axis=-1)
    distance = torch.clamp(torch.nan_to_num(distance), t_samples[:, 0], t_samples[:, -1])
    if white_bkgd:
        comp_rgb = comp_rgb + (1. - torch.unsqueeze(acc, dim=-1))
    return comp_rgb, distance, acc, weights
//  [python]
def get_rotvec(self, eps: float = 1e-6) -> torch.Tensor:
        """Return the underlying axis-angle rotation vector.

        Follow's scipy's implementation:
        https://github.com/scipy/scipy/blob/HEAD/scipy/spatial/transform/_rotation.pyx#L1385-L1402

        Returns:
            The stored rotation as a axis-angle vector.
        """
        quat = self.get_quats()
        # w > 0 to ensure 0 <= angle <= pi
        flip = (quat[..., :1] < 0).float()
        quat = (-1 * quat) * flip + (1 - flip) * quat

        angle = 2 * torch.atan2(torch.norm(quat[..., 1:], dim=-1), quat[..., 0])

        angle2 = angle * angle
        small_angle_scales = 2 + angle2 / 12 + 7 * angle2 * angle2 / 2880
        large_angle_scales = angle / torch.sin(angle / 2 + eps)

        small_angles = (angle <= 1e-3).float()
        rot_vec_scale = small_angle_scales * small_angles + (1 - small_angles) * large_angle_scales
        rot_vec = rot_vec_scale[..., None] * quat[..., 1:]
        return rot_vec
//  [python]
def get_rotvec(self, eps: float = 1e-6) -> torch.Tensor:
        """Return the underlying axis-angle rotation vector.

        Follow's scipy's implementation:
        https://github.com/scipy/scipy/blob/HEAD/scipy/spatial/transform/_rotation.pyx#L1385-L1402

        Returns:
            The stored rotation as a axis-angle vector.
        """
        quat = self.get_quats()
        # w > 0 to ensure 0 <= angle <= pi
        flip = (quat[..., :1] < 0).float()
        quat = (-1 * quat) * flip + (1 - flip) * quat

        angle = 2 * torch.atan2(torch.linalg.norm(quat[..., 1:], dim=-1), quat[..., 0])

        angle2 = angle * angle
        small_angle_scales = 2 + angle2 / 12 + 7 * angle2 * angle2 / 2880
        large_angle_scales = angle / torch.sin(angle / 2 + eps)

        small_angles = (angle <= 1e-3).float()
        rot_vec_scale = small_angle_scales * small_angles + (1 - small_angles) * large_angle_scales
        rot_vec = rot_vec_scale[..., None] * quat[..., 1:]
        return rot_vec
//  [python]
def forward(self, positions):
        # Apply global scale and translation
        positions = (
            positions * self.global_scale[None, None, :]
            + self.global_translation[None, None, :]
        )

        mag = torch.norm(positions, p=self.order, dim=-1)
        mask = mag >= 1
        x_new = positions.clone()
        x_new[mask] = (2 - (1 / mag[mask][..., None])) * (positions[mask] / mag[mask][..., None])

        return x_new
//  [python]
def forward(self, positions):
        # Apply global scale and translation
        positions = (
            positions * self.global_scale[None, None, :]
            + self.global_translation[None, None, :]
        )

        mag = torch.linalg.norm(positions, ord=self.order, dim=-1)
        mask = mag >= 1
        x_new = positions.clone()
        x_new[mask] = (2 - (1 / mag[mask][..., None])) * (positions[mask] / mag[mask][..., None])

        return x_new
//  [python]
def forward(self, x: torch.Tensor) -> torch.Tensor:
        assert (x.shape[-2] == self.input_size) and (
            x.shape[-1] == self.input_size
        ), f"The model expects input images to be resized to {self.input_size}x{self.input_size}"

        res = self.visual.forward(x)
        if self.normalize:
            res = res / torch.norm(res, 2, dim=1, keepdim=True).detach()
        return res
//  [python]
def forward(self, x: torch.Tensor) -> torch.Tensor:
        assert (x.shape[-2] == self.input_size) and (
            x.shape[-1] == self.input_size
        ), f"The model expects input images to be resized to {self.input_size}x{self.input_size}"

        res = self.visual.forward(x)
        if self.normalize:
            res = res / torch.linalg.norm(res, 2, dim=1, keepdim=True).detach()
        return res
//  [python]
def __init__(self, vecX, vecU, beta=1e-4):
        super().__init__()
        N0, N1 = vecX.shape
        self.vecX = torch.tensor(vecX, dtype=torch.float64)
        self.vecU = torch.tensor(vecU, dtype=torch.float64)
        self.vecU_norm = self.vecU / torch.norm(self.vecU)
        np_rng = np.random.default_rng()
        self.beta = beta
        self.lambda_ = torch.nn.Parameter(torch.tensor(np_rng.uniform(0, 1, size=N0), dtype=torch.float64))
        self.alpha = torch.nn.Parameter(torch.tensor([0], dtype=torch.float64))
//  [python]
def __init__(self, vecX, vecU, beta=1e-4):
        super().__init__()
        N0, N1 = vecX.shape
        self.vecX = torch.tensor(vecX, dtype=torch.float64)
        self.vecU = torch.tensor(vecU, dtype=torch.float64)
        self.vecU_norm = self.vecU / torch.linalg.norm(self.vecU)
        np_rng = np.random.default_rng()
        self.beta = beta
        self.lambda_ = torch.nn.Parameter(torch.tensor(np_rng.uniform(0, 1, size=N0), dtype=torch.float64))
        self.alpha = torch.nn.Parameter(torch.tensor([0], dtype=torch.float64))
//  [python]
def forward(self):
        tmp0 = torch.nn.functional.relu(self.lambda_)
        tmp1 = (tmp0 / torch.sum(tmp0)) @ self.vecX
        tmp2 = torch.dot(tmp1, self.vecU_norm) / torch.norm(tmp1)
        loss_angle = torch.sqrt(1-tmp2) #remove nan in gradident
        # loss_angle = torch.nn.functional.smooth_l1_loss(torch.arccos(tmp2), torch.tensor(0.0), beta=self.beta)
        # loss_L1 = torch.abs(self.alpha[0]*self.vecU - tmp1).sum()
        # loss_L1 = torch.nn.functional.smooth_l1_loss(self.alpha[0]*self.vecU, tmp1, beta=self.beta)
        loss_L1 = torch.abs(self.alpha[0]*self.vecU - tmp1).sum()
        loss_alpha = - self.alpha[0]
        return loss_angle, loss_L1, loss_alpha
//  [python]
def forward(self):
        tmp0 = torch.nn.functional.relu(self.lambda_)
        tmp1 = (tmp0 / torch.sum(tmp0)) @ self.vecX
        tmp2 = torch.dot(tmp1, self.vecU_norm) / torch.linalg.norm(tmp1)
        loss_angle = torch.sqrt(1-tmp2) #remove nan in gradident
        # loss_angle = torch.nn.functional.smooth_l1_loss(torch.arccos(tmp2), torch.tensor(0.0), beta=self.beta)
        # loss_L1 = torch.abs(self.alpha[0]*self.vecU - tmp1).sum()
        # loss_L1 = torch.nn.functional.smooth_l1_loss(self.alpha[0]*self.vecU, tmp1, beta=self.beta)
        loss_L1 = torch.abs(self.alpha[0]*self.vecU - tmp1).sum()
        loss_alpha = - self.alpha[0]
        return loss_angle, loss_L1, loss_alpha
//  [python]
def slerp_geometry(control_points, params: interpolation_tensor.InterpolationParams):
    p0, p1 = control_points
    p0_norm = t.norm(p0)
    p1_norm = t.norm(p1)

    similarity = torch.sum((p0 / p0_norm) * (p1 / p1_norm))
    similarity = min(1., max(-1., float(similarity)))
    if similarity <= params.slerp_epsilon - 1 or similarity >= 1 - params.slerp_epsilon:
        return linear_geometry(control_points, params)

    angle = math.acos(float(similarity)) / 2

    slerp_t = angle * (2 * params.t - 1)
    slerp_t = math.tan(slerp_t) / math.tan(angle)
    slerp_t = (slerp_t + 1) / 2

    normalized_p1 = p1 / p1_norm * p0_norm
    slerp_p = p0 + (normalized_p1 - p0) * slerp_t
    slerp_p = slerp_p / t.norm(slerp_p) * (p0_norm + (p1_norm - p0_norm) * params.t)

    lerp_p = linear_geometry(control_points, params)
    return lerp_p + (slerp_p - lerp_p) * params.slerp_scale
//  [python]
def slerp_geometry(control_points, params: interpolation_tensor.InterpolationParams):
    p0, p1 = control_points
    p0_norm = torch.linalg.norm(p0)
    p1_norm = torch.linalg.norm(p1)

    similarity = torch.sum((p0 / p0_norm) * (p1 / p1_norm))
    similarity = min(1., max(-1., float(similarity)))
    if similarity <= params.slerp_epsilon - 1 or similarity >= 1 - params.slerp_epsilon:
        return linear_geometry(control_points, params)

    angle = math.acos(float(similarity)) / 2

    slerp_t = angle * (2 * params.t - 1)
    slerp_t = math.tan(slerp_t) / math.tan(angle)
    slerp_t = (slerp_t + 1) / 2

    normalized_p1 = p1 / p1_norm * p0_norm
    slerp_p = p0 + (normalized_p1 - p0) * slerp_t
    slerp_p = slerp_p / torch.linalg.norm(slerp_p) * (p0_norm + (p1_norm - p0_norm) * params.t)

    lerp_p = linear_geometry(control_points, params)
    return lerp_p + (slerp_p - lerp_p) * params.slerp_scale
//  [python]
def cka_torch(gram_x, gram_y, debiased=False):
    """Compute CKA.

    Args:
      gram_x: A num_examples x num_examples Gram matrix.
      gram_y: A num_examples x num_examples Gram matrix.
      debiased: Use unbiased estimator of HSIC. CKA may still be biased.

    Returns:
      The value of CKA between X and Y.
    """
    gram_x = center_gram_torch(gram_x, unbiased=debiased)
    gram_y = center_gram_torch(gram_y, unbiased=debiased)

    # Note: To obtain HSIC, this should be divided by (n-1)**2 (biased variant) or
    # n*(n-3) (unbiased variant), but this cancels for CKA.
    scaled_hsic = torch.matmul(torch.ravel(gram_x), torch.ravel(gram_y))

    normalization_x = torch.norm(gram_x)
    normalization_y = torch.norm(gram_y)
    return scaled_hsic / (normalization_x * normalization_y)
//  [python]
def cka_torch(gram_x, gram_y, debiased=False):
    """Compute CKA.

    Args:
      gram_x: A num_examples x num_examples Gram matrix.
      gram_y: A num_examples x num_examples Gram matrix.
      debiased: Use unbiased estimator of HSIC. CKA may still be biased.

    Returns:
      The value of CKA between X and Y.
    """
    gram_x = center_gram_torch(gram_x, unbiased=debiased)
    gram_y = center_gram_torch(gram_y, unbiased=debiased)

    # Note: To obtain HSIC, this should be divided by (n-1)**2 (biased variant) or
    # n*(n-3) (unbiased variant), but this cancels for CKA.
    scaled_hsic = torch.matmul(torch.ravel(gram_x), torch.ravel(gram_y))

    normalization_x = torch.linalg.norm(gram_x)
    normalization_y = torch.linalg.norm(gram_y)
    return scaled_hsic / (normalization_x * normalization_y)
//  [python]
def _get_coeff(self):
        pcoeff = self.pcoeff / torch.norm(self.pcoeff)
        qcoeff = self.qcoeff / torch.norm(self.qcoeff)
        qcoeff = torch.complex(qcoeff[:5], 0*qcoeff[5:])
        return pcoeff,qcoeff
//  [python]
def _get_coeff(self):
        pcoeff = self.pcoeff / torch.linalg.norm(self.pcoeff)
        qcoeff = self.qcoeff / torch.linalg.norm(self.qcoeff)
        qcoeff = torch.complex(qcoeff[:5], 0*qcoeff[5:])
        return pcoeff,qcoeff
//  [python]
def torch_find_closest_on_surface(others, refN, pointN, soft_dist_func="softkmin", resource_folder="resource"):
    k = 10
    new_others = torch.empty(others.shape, device=others.device)
    for i in range(refN):
        my_str = resource_folder+"/MNI_templates/xyzall{}.npy".format(str(i + 1))
        xyz = load_raw_MNI_data(my_str, i, resource_folder)
        xyz = torch.FloatTensor(xyz).to(others.device)
        single_instance = others[i].unsqueeze(1)
        distances = torch.norm(single_instance - xyz.repeat(85, 1, 1), dim=-1).double()
        # hard_min = torch.min(distances, dim=-1).values
        # other_soft_min = torch.topk(distances, 1000, dim=-1, largest=False)
        if soft_dist_func == "softkmin":
            x = k_softmin(k, distances)
            p = (x.float() @ xyz)
        else:
            raise NotImplementedError
        # p_real = xyz[torch.argmin(distances, dim=-1)]
        # print("error: ", torch.mean(torch.norm(p - p_real, dim=1)))
        new_others[i] = p
    return torch.mean(new_others, axis=0)
//  [python]
def torch_find_closest_on_surface(others, refN, pointN, soft_dist_func="softkmin", resource_folder="resource"):
    k = 10
    new_others = torch.empty(others.shape, device=others.device)
    for i in range(refN):
        my_str = resource_folder+"/MNI_templates/xyzall{}.npy".format(str(i + 1))
        xyz = load_raw_MNI_data(my_str, i, resource_folder)
        xyz = torch.FloatTensor(xyz).to(others.device)
        single_instance = others[i].unsqueeze(1)
        distances = torch.linalg.norm(single_instance - xyz.repeat(85, 1, 1), dim=-1).double()
        # hard_min = torch.min(distances, dim=-1).values
        # other_soft_min = torch.topk(distances, 1000, dim=-1, largest=False)
        if soft_dist_func == "softkmin":
            x = k_softmin(k, distances)
            p = (x.float() @ xyz)
        else:
            raise NotImplementedError
        # p_real = xyz[torch.argmin(distances, dim=-1)]
        # print("error: ", torch.mean(torch.norm(p - p_real, dim=1)))
        new_others[i] = p
    return torch.mean(new_others, axis=0)
//  [python]
def forward(self):
        coeff = self.theta / torch.norm(self.theta)
        state = (coeff[0]*self.basis[0] + coeff[1]*self.basis[2] + coeff[2]*self.basis[4]).reshape(4,4)
        tmp0 = (state @ state.T) * 3
        tmp1 = (tmp0@tmp0 - tmp0).reshape(-1)
        loss = torch.dot(tmp1,tmp1)
        return loss
//  [python]
def forward(self):
        coeff = self.theta / torch.linalg.norm(self.theta)
        state = (coeff[0]*self.basis[0] + coeff[1]*self.basis[2] + coeff[2]*self.basis[4]).reshape(4,4)
        tmp0 = (state @ state.T) * 3
        tmp1 = (tmp0@tmp0 - tmp0).reshape(-1)
        loss = torch.dot(tmp1,tmp1)
        return loss
//  [python]
def forward(self):
        qcoeff = self.qcoeff / torch.norm(self.qcoeff)
        qstate = (qcoeff @ self.basis).reshape(4,4)
        tmp1 = (self.prdm - qstate @ qstate.T).reshape(-1)
        fidelity = (self.pcoeff[0] * qcoeff[0] + self.pcoeff[1] * qcoeff[2] + self.pcoeff[2] * qcoeff[4])**2
        loss = torch.dot(tmp1, tmp1).real + fidelity
        return loss
//  [python]
def forward(self):
        qcoeff = self.qcoeff / torch.linalg.norm(self.qcoeff)
        qstate = (qcoeff @ self.basis).reshape(4,4)
        tmp1 = (self.prdm - qstate @ qstate.T).reshape(-1)
        fidelity = (self.pcoeff[0] * qcoeff[0] + self.pcoeff[1] * qcoeff[2] + self.pcoeff[2] * qcoeff[4])**2
        loss = torch.dot(tmp1, tmp1).real + fidelity
        return loss
//  [python]
def _test_sqrt_sum_with_cpu_random_data(test_case, placement, sbp):
    x = random_tensor(ndim=4, dim0=8, dim1=32, dim2=40, dim3=64).to_global(
        placement=placement, sbp=sbp
    )
    y = torch.norm(x)
    return y
//  [python]
def _test_sqrt_sum_with_cpu_random_data(test_case, placement, sbp):
    x = random_tensor(ndim=4, dim0=8, dim1=32, dim2=40, dim3=64).to_global(
        placement=placement, sbp=sbp
    )
    y = torch.linalg.norm(x)
    return y
//  [python]
def __init__(self,
        rot_mats: Optional[torch.Tensor] = None,
        quats: Optional[torch.Tensor] = None,
        normalize_quats: bool = True,
    ):
        """
            Args:
                rot_mats:
                    A [*, 3, 3] rotation matrix tensor. Mutually exclusive with
                    quats
                quats:
                    A [*, 4] quaternion. Mutually exclusive with rot_mats. If
                    normalize_quats is not True, must be a unit quaternion
                normalize_quats:
                    If quats is specified, whether to normalize quats
        """
        if((rot_mats is None and quats is None) or 
            (rot_mats is not None and quats is not None)):
            raise ValueError("Exactly one input argument must be specified")

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

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

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def __init__(self,
        rot_mats: Optional[torch.Tensor] = None,
        quats: Optional[torch.Tensor] = None,
        normalize_quats: bool = True,
    ):
        """
            Args:
                rot_mats:
                    A [*, 3, 3] rotation matrix tensor. Mutually exclusive with
                    quats
                quats:
                    A [*, 4] quaternion. Mutually exclusive with rot_mats. If
                    normalize_quats is not True, must be a unit quaternion
                normalize_quats:
                    If quats is specified, whether to normalize quats
        """
        if((rot_mats is None and quats is None) or 
            (rot_mats is not None and quats is not None)):
            raise ValueError("Exactly one input argument must be specified")

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

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

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def _test_scalar_random_data(test_case, placement, sbp):
    x = random_tensor(ndim=4, dim0=8, dim1=24, dim2=16, dim3=40).to_global(
        placement=placement, sbp=sbp
    )
    y = torch.norm(x)
    return y
//  [python]
def _test_scalar_random_data(test_case, placement, sbp):
    x = random_tensor(ndim=4, dim0=8, dim1=24, dim2=16, dim3=40).to_global(
        placement=placement, sbp=sbp
    )
    y = torch.linalg.norm(x)
    return y
//  [python]
def torch_find_closest_on_surface_naive(othersRefList, XYZ, pointN, calc_sd_and_var=False):
    """
    finds closest point on cortical surface for every (transformed) sensor location
    by averaging over 3 closest points on cortical surface
    :param othersRefList: the refN x pointN x 3 transformed sensor locations (into ref brains)
    :param XYZ the raw measurements from template reference brains
    :param pointN: number of sensors
    :return:
    other - location on cortical surface per sensor
    otherVar - variance of each otherH sensor
    otherSD - root of variance of each otherH sensor
    """
    device = othersRefList.device
    other = torch.ones((pointN, 3)).to(device)
    otherVar = torch.ones((pointN, 4)).to(device)
    otherSD = torch.ones((pointN, 4)).to(device)
    top = 3
    for i in range(pointN):
        AA = torch.mean(othersRefList[:, i], dim=0)
        PP = torch.broadcast_to(AA, XYZ.shape)
        D = torch.norm(XYZ - PP, dim=1)
        XYZtop = XYZ[torch.topk(D, largest=False, k=top).indices, :]
        closest = torch.mean(XYZtop, dim=0)
        other[i, :] = closest
        if calc_sd_and_var:
            AAA = othersRefList[:, i]
            AV = closest
            N = AAA.shape[0]
            subMat = torch.ones(AAA.shape).to(device)
            for j in range(N):
                subMat[j, :] = AV

            dispEach = AAA - subMat
            dispEachSq = dispEach * dispEach
            XYZSS = torch.sum(dispEachSq, dim=0)
            RSS = torch.sum(XYZSS)

            XYZVar = XYZSS / (N - 1)
            RSSVar = RSS / (N - 1)

            VV = torch.cat((XYZVar, RSSVar.unsqueeze(0)))
            otherVar[i, :] = VV
            otherSD[i, :] = torch.sqrt(VV)
    return other, otherVar, otherSD
//  [python]
def torch_find_closest_on_surface_naive(othersRefList, XYZ, pointN, calc_sd_and_var=False):
    """
    finds closest point on cortical surface for every (transformed) sensor location
    by averaging over 3 closest points on cortical surface
    :param othersRefList: the refN x pointN x 3 transformed sensor locations (into ref brains)
    :param XYZ the raw measurements from template reference brains
    :param pointN: number of sensors
    :return:
    other - location on cortical surface per sensor
    otherVar - variance of each otherH sensor
    otherSD - root of variance of each otherH sensor
    """
    device = othersRefList.device
    other = torch.ones((pointN, 3)).to(device)
    otherVar = torch.ones((pointN, 4)).to(device)
    otherSD = torch.ones((pointN, 4)).to(device)
    top = 3
    for i in range(pointN):
        AA = torch.mean(othersRefList[:, i], dim=0)
        PP = torch.broadcast_to(AA, XYZ.shape)
        D = torch.linalg.norm(XYZ - PP, dim=1)
        XYZtop = XYZ[torch.topk(D, largest=False, k=top).indices, :]
        closest = torch.mean(XYZtop, dim=0)
        other[i, :] = closest
        if calc_sd_and_var:
            AAA = othersRefList[:, i]
            AV = closest
            N = AAA.shape[0]
            subMat = torch.ones(AAA.shape).to(device)
            for j in range(N):
                subMat[j, :] = AV

            dispEach = AAA - subMat
            dispEachSq = dispEach * dispEach
            XYZSS = torch.sum(dispEachSq, dim=0)
            RSS = torch.sum(XYZSS)

            XYZVar = XYZSS / (N - 1)
            RSSVar = RSS / (N - 1)

            VV = torch.cat((XYZVar, RSSVar.unsqueeze(0)))
            otherVar[i, :] = VV
            otherSD[i, :] = torch.sqrt(VV)
    return other, otherVar, otherSD
//  [python]
def accuracy(pred, y):
    """
    :param pred: predictions
    :param y: ground truth
    :return: accuracy, defined as 1 - (norm(y - pred) / norm(y))
    """
    return 1 - torch.norm(y - pred, "fro") / torch.norm(y, "fro")
//  [python]
def accuracy(pred, y):
    """
    :param pred: predictions
    :param y: ground truth
    :return: accuracy, defined as 1 - (norm(y - pred) / norm(y))
    """
    return 1 - torch.linalg.norm(y - pred, "fro") / torch.linalg.norm(y, "fro")
//  [python]
def full_scene_scale(batch):
    cameras = FoVPerspectiveCameras(R=batch["R"], T=batch["T"], device="cuda")
    cc = cameras.get_camera_center()
    centroid = torch.mean(cc, dim=0)

    diffs = cc - centroid
    norms = torch.norm(diffs, dim=1)

    furthest_index = torch.argmax(norms).item()
    scale = norms[furthest_index].item()
    return scale
//  [python]
def full_scene_scale(batch):
    cameras = FoVPerspectiveCameras(R=batch["R"], T=batch["T"], device="cuda")
    cc = cameras.get_camera_center()
    centroid = torch.mean(cc, dim=0)

    diffs = cc - centroid
    norms = torch.linalg.norm(diffs, dim=1)

    furthest_index = torch.argmax(norms).item()
    scale = norms[furthest_index].item()
    return scale
//  [python]
def __init__(
        self,
        rot_mats: Optional[torch.Tensor] = None,
        quats: Optional[torch.Tensor] = None,
        normalize_quats: bool = True,
    ):
        """
        Args:
            rot_mats:
                A [*, 3, 3] rotation matrix tensor. Mutually exclusive with quats
            quats:
                A [*, 4] quaternion. Mutually exclusive with rot_mats. If normalize_quats is not True, must be a unit
                quaternion
            normalize_quats:
                If quats is specified, whether to normalize quats
        """
        if (rot_mats is None and quats is None) or (
            rot_mats is not None and quats is not None
        ):
            raise ValueError("Exactly one input argument must be specified")

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

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

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

        self._rot_mats = rot_mats
        self._quats = quats