jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_initialize_lengthscale_batch(self):
        kernel = ExponentialDecayKernel(batch_shape=torch.Size([2]))
        ls_init = torch.tensor([1.0, 2.0])
        kernel.initialize(lengthscale=ls_init)
        actual_value = ls_init.view_as(kernel.lengthscale)
        self.assertLess(torch.linalg.norm(kernel.lengthscale - actual_value), 1e-5)
//  [python]
def test_initialize_offset(self):
        kernel = ExponentialDecayKernel()
        kernel.initialize(offset=1)
        actual_value = torch.tensor(1.0).view_as(kernel.offset)
        self.assertLess(torch.norm(kernel.offset - actual_value), 1e-5)
//  [python]
def test_initialize_offset(self):
        kernel = ExponentialDecayKernel()
        kernel.initialize(offset=1)
        actual_value = torch.tensor(1.0).view_as(kernel.offset)
        self.assertLess(torch.linalg.norm(kernel.offset - actual_value), 1e-5)
//  [python]
def test_initialize_offset_batch(self):
        kernel = ExponentialDecayKernel(batch_shape=torch.Size([2]))
        off_init = torch.tensor([1.0, 2.0])
        kernel.initialize(offset=off_init)
        actual_value = off_init.view_as(kernel.offset)
        self.assertLess(torch.norm(kernel.offset - actual_value), 1e-5)
//  [python]
def test_initialize_offset_batch(self):
        kernel = ExponentialDecayKernel(batch_shape=torch.Size([2]))
        off_init = torch.tensor([1.0, 2.0])
        kernel.initialize(offset=off_init)
        actual_value = off_init.view_as(kernel.offset)
        self.assertLess(torch.linalg.norm(kernel.offset - actual_value), 1e-5)
//  [python]
def test_initialize_power(self):
        kernel = ExponentialDecayKernel()
        kernel.initialize(power=1)
        actual_value = torch.tensor(1.0).view_as(kernel.power)
        self.assertLess(torch.norm(kernel.power - actual_value), 1e-5)
//  [python]
def test_initialize_power(self):
        kernel = ExponentialDecayKernel()
        kernel.initialize(power=1)
        actual_value = torch.tensor(1.0).view_as(kernel.power)
        self.assertLess(torch.linalg.norm(kernel.power - actual_value), 1e-5)
//  [python]
def test_initialize_power_batch(self):
        kernel = ExponentialDecayKernel(batch_shape=torch.Size([2]))
        power_init = torch.tensor([1.0, 2.0])
        kernel.initialize(power=power_init)
        actual_value = power_init.view_as(kernel.power)
        self.assertLess(torch.norm(kernel.power - actual_value), 1e-5)
//  [python]
def test_initialize_power_batch(self):
        kernel = ExponentialDecayKernel(batch_shape=torch.Size([2]))
        power_init = torch.tensor([1.0, 2.0])
        kernel.initialize(power=power_init)
        actual_value = power_init.view_as(kernel.power)
        self.assertLess(torch.linalg.norm(kernel.power - actual_value), 1e-5)
//  [python]
def _check_tensor_impl(tensor: torch.Tensor) -> bool:
        with torch.no_grad():
            if tensor.ndim != 2 or tensor.shape[1] != 2:
                raise ValueError("SO2 data tensors can only be 2D vectors.")

            MATRIX_EPS = _THESEUS_GLOBAL_PARAMS.get_eps("so2", "matrix", tensor.dtype)
            if tensor.dtype != torch.float64:
                tensor = tensor.double()

            _check = (
                tensor.norm(dim=1) - 1
            ).abs().max().item() <= MATRIX_EPS

        return _check
//  [python]
def _check_tensor_impl(tensor: torch.Tensor) -> bool:
        with torch.no_grad():
            if tensor.ndim != 2 or tensor.shape[1] != 2:
                raise ValueError("SO2 data tensors can only be 2D vectors.")

            MATRIX_EPS = _THESEUS_GLOBAL_PARAMS.get_eps("so2", "matrix", tensor.dtype)
            if tensor.dtype != torch.float64:
                tensor = tensor.double()

            _check = (
                torch.linalg.norm(tensor, dim=1) - 1
            ).abs().max().item() <= MATRIX_EPS

        return _check
//  [python]
def cka(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(gram_x, unbiased=debiased)
    gram_y = center_gram(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 = gram_x.view(-1) @ gram_y.view(-1)

    normalization_x = torch.norm(gram_x)
    normalization_y = torch.norm(gram_y)
    return scaled_hsic / (normalization_x * normalization_y)
//  [python]
def cka(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(gram_x, unbiased=debiased)
    gram_y = center_gram(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 = gram_x.view(-1) @ gram_y.view(-1)

    normalization_x = torch.linalg.norm(gram_x)
    normalization_y = torch.linalg.norm(gram_y)
    return scaled_hsic / (normalization_x * normalization_y)
//  [python]
def avg_R_from_2poses(hands1, hands2, valids1, valids2, is_torch=False, quatn_mo=None, thres=0.4):
    Rs = []
    for i in range(hands1.shape[0]):
        pred1, pred2 = hands1[i], hands2[i]
        p1, p2 = [], []
        for j in range(pred1.shape[0]):
            if valids1[i, j] and valids2[i, j]:
                p1.append(pred1[j])
                p2.append(pred2[j])
        Rs.append(R_from_2poses(p1, p2, is_torch=is_torch))

    quatns = [quaternion_from_matrix(R, is_torch=is_torch) for R in Rs]
    if is_torch:
        quatn_avg_raw = torch.mean(torch.stack(quatns), dim=0)
        quatn_avg = quatn_avg_raw / torch.norm(quatn_avg_raw)
    else:
        quatn_avg = np.mean(quatns, axis=0)
        quatn_avg /= np.linalg.norm(quatn_avg)
    R_avg = matrix_from_quaternion(quatn_avg, is_torch=is_torch)

    return R_avg, quatn_avg
//  [python]
def avg_R_from_2poses(hands1, hands2, valids1, valids2, is_torch=False, quatn_mo=None, thres=0.4):
    Rs = []
    for i in range(hands1.shape[0]):
        pred1, pred2 = hands1[i], hands2[i]
        p1, p2 = [], []
        for j in range(pred1.shape[0]):
            if valids1[i, j] and valids2[i, j]:
                p1.append(pred1[j])
                p2.append(pred2[j])
        Rs.append(R_from_2poses(p1, p2, is_torch=is_torch))

    quatns = [quaternion_from_matrix(R, is_torch=is_torch) for R in Rs]
    if is_torch:
        quatn_avg_raw = torch.mean(torch.stack(quatns), dim=0)
        quatn_avg = quatn_avg_raw / torch.linalg.norm(quatn_avg_raw)
    else:
        quatn_avg = np.mean(quatns, axis=0)
        quatn_avg /= np.linalg.norm(quatn_avg)
    R_avg = matrix_from_quaternion(quatn_avg, is_torch=is_torch)

    return R_avg, quatn_avg
//  [python]
def make_rotX(angs, eps=1e-6):
    B,L = angs.shape[:2]
    NORM = torch.norm(angs, dim=-1) + eps

    RTs = torch.eye(4,  device=angs.device).repeat(B,L,1,1)

    RTs[:,:,1,1] = angs[:,:,0]/NORM
    RTs[:,:,1,2] = -angs[:,:,1]/NORM
    RTs[:,:,2,1] = angs[:,:,1]/NORM
    RTs[:,:,2,2] = angs[:,:,0]/NORM
    return RTs
//  [python]
def make_rotX(angs, eps=1e-6):
    B,L = angs.shape[:2]
    NORM = torch.linalg.norm(angs, dim=-1) + eps

    RTs = torch.eye(4,  device=angs.device).repeat(B,L,1,1)

    RTs[:,:,1,1] = angs[:,:,0]/NORM
    RTs[:,:,1,2] = -angs[:,:,1]/NORM
    RTs[:,:,2,1] = angs[:,:,1]/NORM
    RTs[:,:,2,2] = angs[:,:,0]/NORM
    return RTs
//  [python]
def matrix_from_quaternion(q, is_torch=False):
    if is_torch:
        R = torch.empty((3, 3), dtype=q.dtype, device=q.device)
        q = q / torch.norm(q)
    else:
        R = np.empty((3, 3))
        q = q / np.linalg.norm(q)  # Normalize the quaternion if necessary

    q0, q1, q2, q3 = q[0], q[1], q[2], q[3]
    q0q0, q1q1, q2q2, q3q3 = q0 * q0, q1 * q1, q2 * q2, q3 * q3
    q0q1, q0q2, q0q3 = q0 * q1, q0 * q2, q0 * q3
    q1q2, q1q3, q2q3 = q1 * q2, q1 * q3, q2 * q3

    R[0, 0] = q0q0 + q1q1 - q2q2 - q3q3
    R[0, 1] = 2 * (q1q2 - q0q3)
    R[0, 2] = 2 * (q1q3 + q0q2)

    R[1, 0] = 2 * (q1q2 + q0q3)
    R[1, 1] = q0q0 - q1q1 + q2q2 - q3q3
    R[1, 2] = 2 * (q2q3 - q0q1)

    R[2, 0] = 2 * (q1q3 - q0q2)
    R[2, 1] = 2 * (q2q3 + q0q1)
    R[2, 2] = q0q0 - q1q1 - q2q2 + q3q3

    return R
//  [python]
def matrix_from_quaternion(q, is_torch=False):
    if is_torch:
        R = torch.empty((3, 3), dtype=q.dtype, device=q.device)
        q = q / torch.linalg.norm(q)
    else:
        R = np.empty((3, 3))
        q = q / np.linalg.norm(q)  # Normalize the quaternion if necessary

    q0, q1, q2, q3 = q[0], q[1], q[2], q[3]
    q0q0, q1q1, q2q2, q3q3 = q0 * q0, q1 * q1, q2 * q2, q3 * q3
    q0q1, q0q2, q0q3 = q0 * q1, q0 * q2, q0 * q3
    q1q2, q1q3, q2q3 = q1 * q2, q1 * q3, q2 * q3

    R[0, 0] = q0q0 + q1q1 - q2q2 - q3q3
    R[0, 1] = 2 * (q1q2 - q0q3)
    R[0, 2] = 2 * (q1q3 + q0q2)

    R[1, 0] = 2 * (q1q2 + q0q3)
    R[1, 1] = q0q0 - q1q1 + q2q2 - q3q3
    R[1, 2] = 2 * (q2q3 - q0q1)

    R[2, 0] = 2 * (q1q3 - q0q2)
    R[2, 1] = 2 * (q2q3 + q0q1)
    R[2, 2] = q0q0 - q1q1 - q2q2 + q3q3

    return R
//  [python]
def make_rotZ(angs, eps=1e-6):
    B,L = angs.shape[:2]
    NORM = torch.norm(angs, dim=-1) + eps

    RTs = torch.eye(4,  device=angs.device).repeat(B,L,1,1)

    RTs[:,:,0,0] = angs[:,:,0]/NORM
    RTs[:,:,0,1] = -angs[:,:,1]/NORM
    RTs[:,:,1,0] = angs[:,:,1]/NORM
    RTs[:,:,1,1] = angs[:,:,0]/NORM
    return RTs
//  [python]
def make_rotZ(angs, eps=1e-6):
    B,L = angs.shape[:2]
    NORM = torch.linalg.norm(angs, dim=-1) + eps

    RTs = torch.eye(4,  device=angs.device).repeat(B,L,1,1)

    RTs[:,:,0,0] = angs[:,:,0]/NORM
    RTs[:,:,0,1] = -angs[:,:,1]/NORM
    RTs[:,:,1,0] = angs[:,:,1]/NORM
    RTs[:,:,1,1] = angs[:,:,0]/NORM
    return RTs
//  [python]
def random_quat(shape: tuple[int] | int, device: Optional[str] = None) -> Tensor:
    """
    Uniformly samples SO3 elements parameterized as quaternions.

    Arguments:
        - shape: Int or tuple denoting the shape of the output tensor.
        - device: device on which the new tensor is created.

    Returns:
        - Tensor of shape `(*shape, 4)`.
    """
    shape = shape if type(shape) is tuple else (shape,)

    q = torch.randn(*shape, 4, device=device)
    q = q / torch.norm(q, dim=-1, keepdim=True)

    return q
//  [python]
def random_quat(shape: tuple[int] | int, device: Optional[str] = None) -> Tensor:
    """
    Uniformly samples SO3 elements parameterized as quaternions.

    Arguments:
        - shape: Int or tuple denoting the shape of the output tensor.
        - device: device on which the new tensor is created.

    Returns:
        - Tensor of shape `(*shape, 4)`.
    """
    shape = shape if type(shape) is tuple else (shape,)

    q = torch.randn(*shape, 4, device=device)
    q = q / torch.linalg.norm(q, dim=-1, keepdim=True)

    return q
//  [python]
def feature_space_linear_cka(features_x, features_y, debiased=False):
    """Compute CKA with a linear kernel, in feature space.

    This is typically faster than computing the Gram matrix when there are fewer
    features than examples.

    Args:
        features_x: A num_examples x num_features matrix of features.
        features_y: A num_examples x num_features matrix of features.
        debiased: Use unbiased estimator of dot product similarity. CKA may still be
        biased. Note that this estimator may be negative.

    Returns:
        The value of CKA between X and Y.
    """
    features_x = features_x - torch.mean(features_x, 0, keepdims=True)
    features_y = features_y - torch.mean(features_y, 0, keepdims=True)

    dot_product_similarity = torch.norm(features_x.T @ (features_y)) ** 2
    normalization_x = torch.norm(features_x.T @ (features_x))
    normalization_y = torch.norm(features_y.T @ (features_y))

    if debiased:
        n = features_x.size(0)
        # Equivalent to np.sum(features_x ** 2, 1) but avoids an intermediate array.
        sum_squared_rows_x = torch.einsum("ij,ij->i", features_x, features_x)
        sum_squared_rows_y = torch.einsum("ij,ij->i", features_y, features_y)
        squared_norm_x = torch.sum(sum_squared_rows_x)
        squared_norm_y = torch.sum(sum_squared_rows_y)

        dot_product_similarity = _debiased_dot_product_similarity_helper(
            dot_product_similarity, sum_squared_rows_x, sum_squared_rows_y, squared_norm_x, squared_norm_y, n
        )
        normalization_x = torch.sqrt(
            _debiased_dot_product_similarity_helper(
                normalization_x**2, sum_squared_rows_x, sum_squared_rows_x, squared_norm_x, squared_norm_x, n
            )
        )
        normalization_y = torch.sqrt(
            _debiased_dot_product_similarity_helper(
                normalization_y**2, sum_squared_rows_y, sum_squared_rows_y, squared_norm_y, squared_norm_y, n
            )
        )

    return dot_product_similarity / (normalization_x * normalization_y)
//  [python]
def feature_space_linear_cka(features_x, features_y, debiased=False):
    """Compute CKA with a linear kernel, in feature space.

    This is typically faster than computing the Gram matrix when there are fewer
    features than examples.

    Args:
        features_x: A num_examples x num_features matrix of features.
        features_y: A num_examples x num_features matrix of features.
        debiased: Use unbiased estimator of dot product similarity. CKA may still be
        biased. Note that this estimator may be negative.

    Returns:
        The value of CKA between X and Y.
    """
    features_x = features_x - torch.mean(features_x, 0, keepdims=True)
    features_y = features_y - torch.mean(features_y, 0, keepdims=True)

    dot_product_similarity = torch.linalg.norm(features_x.T @ (features_y)) ** 2
    normalization_x = torch.linalg.norm(features_x.T @ (features_x))
    normalization_y = torch.linalg.norm(features_y.T @ (features_y))

    if debiased:
        n = features_x.size(0)
        # Equivalent to np.sum(features_x ** 2, 1) but avoids an intermediate array.
        sum_squared_rows_x = torch.einsum("ij,ij->i", features_x, features_x)
        sum_squared_rows_y = torch.einsum("ij,ij->i", features_y, features_y)
        squared_norm_x = torch.sum(sum_squared_rows_x)
        squared_norm_y = torch.sum(sum_squared_rows_y)

        dot_product_similarity = _debiased_dot_product_similarity_helper(
            dot_product_similarity, sum_squared_rows_x, sum_squared_rows_y, squared_norm_x, squared_norm_y, n
        )
        normalization_x = torch.sqrt(
            _debiased_dot_product_similarity_helper(
                normalization_x**2, sum_squared_rows_x, sum_squared_rows_x, squared_norm_x, squared_norm_x, n
            )
        )
        normalization_y = torch.sqrt(
            _debiased_dot_product_similarity_helper(
                normalization_y**2, sum_squared_rows_y, sum_squared_rows_y, squared_norm_y, squared_norm_y, n
            )
        )

    return dot_product_similarity / (normalization_x * normalization_y)
//  [python]
def calculate(self, data_batch, label_batch):
        cove_set = set()
        SA_batch = []
        batch_size = label_batch.size(0)
        layer_output_dict = tool.get_layer_output(self.model, data_batch)
        for (layer_name, layer_output) in layer_output_dict.items():
            SA_batch.append(layer_output[:, self.mask_index_dict[layer_name]].view(batch_size, -1))
        SA_batch = torch.cat(SA_batch, 1).detach().cpu().numpy() # [batch_size, num_neuron]
        for i, label in enumerate(label_batch):
            SA = SA_batch[i]
            
            # # using numpy 
            # dist_a_list = np.linalg.norm(SA - self.SA_cache[int(label.cpu())], axis=1)
            # idx_a = np.argmin(dist_a_list, 0)

            dist_a_list = torch.norm(
                torch.from_numpy(SA).to(self.device) - torch.from_numpy(self.SA_cache[int(label.cpu())]).to(self.device), dim=1)
            idx_a = torch.argmin(dist_a_list, 0).item()

            (SA_a, dist_a) = (self.SA_cache[int(label.cpu())][idx_a], dist_a_list[idx_a])
            dist_a = dist_a.cpu().numpy()

            dist_b_list = []
            for j in range(self.num_class):
                if ( j != int(label.cpu()) ) and ( j in self.SA_cache.keys() ):
                    # # using numpy
                    # dist_b_list += np.linalg.norm(SA - self.SA_cache[j], axis=1).tolist()
                    dist_b_list += torch.norm(
                        torch.from_numpy(SA).to(self.device) - torch.from_numpy(self.SA_cache[j]).to(self.device),
                        dim=1).cpu().numpy().tolist()
            
            dist_b = np.min(dist_b_list)
            dsa = dist_a / dist_b if dist_b > 0 else 1e-6
            if (not np.isnan(dsa)) and (not np.isinf(dsa)):
                cove_set.add(int(dsa / self.threshold))
        cove_set = self.coverage_set.union(cove_set)
        return cove_set
//  [python]
def calculate(self, data_batch, label_batch):
        cove_set = set()
        SA_batch = []
        batch_size = label_batch.size(0)
        layer_output_dict = tool.get_layer_output(self.model, data_batch)
        for (layer_name, layer_output) in layer_output_dict.items():
            SA_batch.append(layer_output[:, self.mask_index_dict[layer_name]].view(batch_size, -1))
        SA_batch = torch.cat(SA_batch, 1).detach().cpu().numpy() # [batch_size, num_neuron]
        for i, label in enumerate(label_batch):
            SA = SA_batch[i]
            
            # # using numpy 
            # dist_a_list = np.linalg.norm(SA - self.SA_cache[int(label.cpu())], axis=1)
            # idx_a = np.argmin(dist_a_list, 0)

            dist_a_list = torch.linalg.norm(
                torch.from_numpy(SA).to(self.device) - torch.from_numpy(self.SA_cache[int(label.cpu())]).to(self.device), dim=1)
            idx_a = torch.argmin(dist_a_list, 0).item()

            (SA_a, dist_a) = (self.SA_cache[int(label.cpu())][idx_a], dist_a_list[idx_a])
            dist_a = dist_a.cpu().numpy()

            dist_b_list = []
            for j in range(self.num_class):
                if ( j != int(label.cpu()) ) and ( j in self.SA_cache.keys() ):
                    # # using numpy
                    # dist_b_list += np.linalg.norm(SA - self.SA_cache[j], axis=1).tolist()
                    dist_b_list += torch.linalg.norm(
                        torch.from_numpy(SA).to(self.device) - torch.from_numpy(self.SA_cache[j]).to(self.device),
                        dim=1).cpu().numpy().tolist()
            
            dist_b = np.min(dist_b_list)
            dsa = dist_a / dist_b if dist_b > 0 else 1e-6
            if (not np.isnan(dsa)) and (not np.isinf(dsa)):
                cove_set.add(int(dsa / self.threshold))
        cove_set = self.coverage_set.union(cove_set)
        return cove_set
//  [python]
def quaternion_log(q: Tensor) -> Tensor:
    """
    Calculates logarithm of quaternions.

    NOTE: log of identity quaternion will be set to the zero
          vector.

    Arguments:
        - q: Tensor of shape `(N, 4)`.

    Returns:
        - Tensor of shape `(N, 4)`.
    """
    q_norm = torch.norm(q, dim=-1, keepdim=True)

    inverse_v_norm = 1 / torch.norm(q[:, 1:], dim=-1, keepdim=True)
    inverse_v_norm[inverse_v_norm == torch.inf] = 0

    log_q_norm = torch.log(q_norm)
    vector_part = (inverse_v_norm * torch.acos(q[:, 0].view(-1, 1) / q_norm)) * q[:, 1:]

    return torch.hstack((log_q_norm, vector_part))
//  [python]
def quaternion_log(q: Tensor) -> Tensor:
    """
    Calculates logarithm of quaternions.

    NOTE: log of identity quaternion will be set to the zero
          vector.

    Arguments:
        - q: Tensor of shape `(N, 4)`.

    Returns:
        - Tensor of shape `(N, 4)`.
    """
    q_norm = torch.linalg.norm(q, dim=-1, keepdim=True)

    inverse_v_norm = 1 / torch.linalg.norm(q[:, 1:], dim=-1, keepdim=True)
    inverse_v_norm[inverse_v_norm == torch.inf] = 0

    log_q_norm = torch.log(q_norm)
    vector_part = (inverse_v_norm * torch.acos(q[:, 0].view(-1, 1) / q_norm)) * q[:, 1:]

    return torch.hstack((log_q_norm, vector_part))
//  [python]
def make_rot_axis(angs, u, eps=1e-6):
    B,L = angs.shape[:2]
    NORM = torch.norm(angs, dim=-1) + eps

    RTs = torch.eye(4,  device=angs.device).repeat(B,L,1,1)

    ct = angs[:,:,0]/NORM
    st = angs[:,:,1]/NORM
    u0 = u[:,:,0]
    u1 = u[:,:,1]
    u2 = u[:,:,2]

    RTs[:,:,0,0] = ct+u0*u0*(1-ct)
    RTs[:,:,0,1] = u0*u1*(1-ct)-u2*st
    RTs[:,:,0,2] = u0*u2*(1-ct)+u1*st
    RTs[:,:,1,0] = u0*u1*(1-ct)+u2*st
    RTs[:,:,1,1] = ct+u1*u1*(1-ct)
    RTs[:,:,1,2] = u1*u2*(1-ct)-u0*st
    RTs[:,:,2,0] = u0*u2*(1-ct)-u1*st
    RTs[:,:,2,1] = u1*u2*(1-ct)+u0*st
    RTs[:,:,2,2] = ct+u2*u2*(1-ct)
    return RTs
//  [python]
def make_rot_axis(angs, u, eps=1e-6):
    B,L = angs.shape[:2]
    NORM = torch.linalg.norm(angs, dim=-1) + eps

    RTs = torch.eye(4,  device=angs.device).repeat(B,L,1,1)

    ct = angs[:,:,0]/NORM
    st = angs[:,:,1]/NORM
    u0 = u[:,:,0]
    u1 = u[:,:,1]
    u2 = u[:,:,2]

    RTs[:,:,0,0] = ct+u0*u0*(1-ct)
    RTs[:,:,0,1] = u0*u1*(1-ct)-u2*st
    RTs[:,:,0,2] = u0*u2*(1-ct)+u1*st
    RTs[:,:,1,0] = u0*u1*(1-ct)+u2*st
    RTs[:,:,1,1] = ct+u1*u1*(1-ct)
    RTs[:,:,1,2] = u1*u2*(1-ct)-u0*st
    RTs[:,:,2,0] = u0*u2*(1-ct)-u1*st
    RTs[:,:,2,1] = u1*u2*(1-ct)+u0*st
    RTs[:,:,2,2] = ct+u2*u2*(1-ct)
    return RTs
//  [python]
def random_s2(shape: tuple[int, ...], device: Optional[str] = None) -> torch.Tensor:
    """
    Generates Tensor of uniformly sampled spherical coordinates on
    S2.

    Arguments:
        - shape: Shape of the output tensor.
        - device: Device on which the new tensor is created.

    Returns:
        - Tensor of shape (*shape, 3).
    """
    x = torch.randn((*shape, 3), device=device)
    return euclid_to_spherical(x / torch.norm(x, dim=-1, keepdim=True))
//  [python]
def random_s2(shape: tuple[int, ...], device: Optional[str] = None) -> torch.Tensor:
    """
    Generates Tensor of uniformly sampled spherical coordinates on
    S2.

    Arguments:
        - shape: Shape of the output tensor.
        - device: Device on which the new tensor is created.

    Returns:
        - Tensor of shape (*shape, 3).
    """
    x = torch.randn((*shape, 3), device=device)
    return euclid_to_spherical(x / torch.linalg.norm(x, dim=-1, keepdim=True))
//  [python]
def QUEscore(temp_feats, n_dim):

    n_samples = temp_feats.shape[1]
    alpha = 4.0
    Sigma = torch.matmul(temp_feats, temp_feats.T) / n_samples
    I = torch.eye(n_dim).cuda()
    Q = torch.exp((alpha * (Sigma - I)) / (torch.norm(Sigma) - 1))
    trace_Q = torch.trace(Q)

    taus = []
    for i in range(n_samples):
        h_i = temp_feats[:, i:i + 1]
        tau_i = torch.matmul(h_i.T, torch.matmul(Q, h_i)) / trace_Q
        tau_i = tau_i.item()
        taus.append(tau_i)
    taus = np.array(taus)

    return taus
//  [python]
def QUEscore(temp_feats, n_dim):

    n_samples = temp_feats.shape[1]
    alpha = 4.0
    Sigma = torch.matmul(temp_feats, temp_feats.T) / n_samples
    I = torch.eye(n_dim).cuda()
    Q = torch.exp((alpha * (Sigma - I)) / (torch.linalg.norm(Sigma, ord=2) - 1))
    trace_Q = torch.trace(Q)

    taus = []
    for i in range(n_samples):
        h_i = temp_feats[:, i:i + 1]
        tau_i = torch.matmul(h_i.T, torch.matmul(Q, h_i)) / trace_Q
        tau_i = tau_i.item()
        taus.append(tau_i)
    taus = np.array(taus)

    return taus
//  [python]
def relative_translation_error(gt_translations, translations):
    r"""Isotropic Relative Rotation Error.

    RTE = \lVert t - \bar{t} \rVert_2

    Args:
        gt_translations (Tensor): ground truth translation vector (*, 3)
        translations (Tensor): estimated translation vector (*, 3)

    Returns:
        rre (Tensor): relative rotation errors (*)
    """
    rte = torch.norm(gt_translations - translations, dim=-1)
    return rte
//  [python]
def relative_translation_error(gt_translations, translations):
    r"""Isotropic Relative Rotation Error.

    RTE = \lVert t - \bar{t} \rVert_2

    Args:
        gt_translations (Tensor): ground truth translation vector (*, 3)
        translations (Tensor): estimated translation vector (*, 3)

    Returns:
        rre (Tensor): relative rotation errors (*)
    """
    rte = torch.linalg.norm(gt_translations - translations, dim=-1)
    return rte
//  [python]
def get_costs(problems, tours):
    # Check that tours are valid, i.e. contain 0 to n -1
    assert (
        torch.arange(tours.size(1), out=tours.data.new()).view(1, -1).expand_as(tours) == tours.data.sort(1)[0]
    ).all(), "Invalid tour"

    # Gather city coordinates in order of tour
    coords = problems.gather(1, tours.unsqueeze(-1).expand_as(problems))

    # Length is distance (L2-norm of difference) from each next location from its prev and of last from first
    rolled_coords = coords.roll(-1, dims=1)
    tour_lens = torch.norm(coords - rolled_coords, dim=2).sum(1)
    return tour_lens
//  [python]
def get_costs(problems, tours):
    # Check that tours are valid, i.e. contain 0 to n -1
    assert (
        torch.arange(tours.size(1), out=tours.data.new()).view(1, -1).expand_as(tours) == tours.data.sort(1)[0]
    ).all(), "Invalid tour"

    # Gather city coordinates in order of tour
    coords = problems.gather(1, tours.unsqueeze(-1).expand_as(problems))

    # Length is distance (L2-norm of difference) from each next location from its prev and of last from first
    rolled_coords = coords.roll(-1, dims=1)
    tour_lens = torch.linalg.norm(coords - rolled_coords, dim=2).sum(1)
    return tour_lens
//  [python]
def forward(self, X):
        # we take the norm and sum over the q-batch dim
        if len(X.shape) > 2:
            return torch.norm(X, dim=-1).sum(-1)
        else:
            return torch.norm(X, dim=-1).squeeze(-1)
//  [python]
def forward(self, X):
        # we take the norm and sum over the q-batch dim
        if len(X.shape) > 2:
            return torch.linalg.norm(X, dim=-1).sum(-1)
        else:
            return torch.linalg.norm(X, dim=-1).squeeze(-1)
//  [python]
def embed_images(self, xs: Iterable[Optional[ImageType]]) -> torch.Tensor:
        """
        :param xs: N images, stored as numpy arrays, tensors, or PIL images.
        :return: an [N x D] tensor of features.
        """
        clip_inputs = self.images_to_tensor(xs)
        results = self.clip_model.encode_image(clip_inputs).float()
        return results / torch.norm(results, dim=-1, keepdim=True)
//  [python]
def embed_images(self, xs: Iterable[Optional[ImageType]]) -> torch.Tensor:
        """
        :param xs: N images, stored as numpy arrays, tensors, or PIL images.
        :return: an [N x D] tensor of features.
        """
        clip_inputs = self.images_to_tensor(xs)
        results = self.clip_model.encode_image(clip_inputs).float()
        return results / torch.linalg.norm(results, dim=-1, keepdim=True)
//  [python]
def embed_text(self, prompts: Iterable[str]) -> torch.Tensor:
        """
        Embed text prompts as an [N x D] tensor.
        """
        enc = self.clip_model.encode_text(
            self._tokenize(list(prompts), truncate=True).to(self.device)
        ).float()
        return enc / torch.norm(enc, dim=-1, keepdim=True)
//  [python]
def embed_text(self, prompts: Iterable[str]) -> torch.Tensor:
        """
        Embed text prompts as an [N x D] tensor.
        """
        enc = self.clip_model.encode_text(
            self._tokenize(list(prompts), truncate=True).to(self.device)
        ).float()
        return enc / torch.linalg.norm(enc, dim=-1, keepdim=True)
//  [python]
def compute_stable_rank(tensor):
    # Compute the stable rank of a tensor
    tensor = tensor.detach()
    fro_norm = torch.norm(tensor)**2
    l2_norm = torch.norm(tensor, ord=2)**2
    rank = fro_norm / l2_norm
    rank = rank
    return rank
//  [python]
def compute_stable_rank(tensor):
    # Compute the stable rank of a tensor
    tensor = tensor.detach()
    fro_norm = torch.linalg.norm(tensor, ord='fro')**2
    l2_norm = torch.linalg.norm(tensor, ord=2)**2
    rank = fro_norm / l2_norm
    rank = rank
    return rank
//  [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.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, exact_dtype=False)

                res_out = torch.tensor([]).to(device)
                torch.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, exact_dtype=False)

            # matrix norm
            x = torch.randn(25, 25, device=device, dtype=dtype)
            xn = x.cpu().numpy()
            for ord in matrix_ords:
                res = torch.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, exact_dtype=False)

                res_out = torch.tensor([]).to(device)
                torch.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, exact_dtype=False)
//  [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, exact_dtype=False)

                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, exact_dtype=False)

            # 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, exact_dtype=False)

                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, exact_dtype=False)
//  [python]
def test_norm_extreme_values(self, device):
        vector_ords = [0, 1, 2, 3, inf, -1, -2, -3, -inf]
        # matrix_ords 'nuc', 2, -2 are skipped currently
        # See issue https://github.com/pytorch/pytorch/issues/71911
        matrix_ords = ['fro', 1, inf, -1, -inf]
        vectors = []
        matrices = []
        for pair in itertools.product([inf, -inf, 0.0, nan, 1.0], repeat=2):
            vectors.append(list(pair))
            matrices.append([[pair[0], pair[1]]])
            matrices.append([[pair[0]], [pair[1]]])
        for vector in vectors:
            x = torch.tensor(vector).to(device)
            x_n = x.cpu().numpy()
            for ord in vector_ords:
                msg = f'ord={ord}, vector={vector}'
                result = torch.norm(x, p=ord)
                result_n = np.linalg.norm(x_n, ord=ord)
                self.assertEqual(result, result_n, msg=msg)

        # TODO: Remove this function once the broken cases are fixed
        def is_broken_matrix_norm_case(ord, x):
            if self.device_type == 'cuda':
                if x.size() == torch.Size([1, 2]):
                    if ord in ['nuc', 2, -2] and isnan(x[0][0]) and x[0][1] == 1:
                        # These cases are broken because of an issue with svd
                        # https://github.com/pytorch/pytorch/issues/43567
                        return True
                if ord in ['nuc', 2, -2]:
                    # These cases are broken because of another issue with svd
                    # https://github.com/pytorch/pytorch/issues/52633
                    return True
            return False

        for matrix in matrices:
            x = torch.tensor(matrix).to(device)
            x_n = x.cpu().numpy()
            for ord in matrix_ords:
                msg = f'ord={ord}, matrix={matrix}'
                if is_broken_matrix_norm_case(ord, x):
                    continue
                else:
                    result_n = np.linalg.norm(x_n, ord=ord)
                    result = torch.norm(x, p=ord)
                    self.assertEqual(result, result_n, msg=msg)
//  [python]
def test_norm_extreme_values(self, device):
        vector_ords = [0, 1, 2, 3, inf, -1, -2, -3, -inf]
        # matrix_ords 'nuc', 2, -2 are skipped currently
        # See issue https://github.com/pytorch/pytorch/issues/71911
        matrix_ords = ['fro', 1, inf, -1, -inf]
        vectors = []
        matrices = []
        for pair in itertools.product([inf, -inf, 0.0, nan, 1.0], repeat=2):
            vectors.append(list(pair))
            matrices.append([[pair[0], pair[1]]])
            matrices.append([[pair[0]], [pair[1]]])
        for vector in vectors:
            x = torch.tensor(vector).to(device)
            x_n = x.cpu().numpy()
            for ord in vector_ords:
                msg = f'ord={ord}, vector={vector}'
                result = torch.linalg.norm(x, ord=ord)
                result_n = np.linalg.norm(x_n, ord=ord)
                self.assertEqual(result, result_n, msg=msg)

        # TODO: Remove this function once the broken cases are fixed
        def is_broken_matrix_norm_case(ord, x):
            if self.device_type == 'cuda':
                if x.size() == torch.Size([1, 2]):
                    if ord in ['nuc', 2, -2] and isnan(x[0][0]) and x[0][1] == 1:
                        # These cases are broken because of an issue with svd
                        # https://github.com/pytorch/pytorch/issues/43567
                        return True
                if ord in ['nuc', 2, -2]:
                    # These cases are broken because of another issue with svd
                    # https://github.com/pytorch/pytorch/issues/52633
                    return True
            return False

        for matrix in matrices:
            x = torch.tensor(matrix).to(device)
            x_n = x.cpu().numpy()
            for ord in matrix_ords:
                msg = f'ord={ord}, matrix={matrix}'
                if is_broken_matrix_norm_case(ord, x):
                    continue
                else:
                    result_n = np.linalg.norm(x_n, ord=ord)
                    result = torch.linalg.norm(x, ord=ord)
                    self.assertEqual(result, result_n, msg=msg)
//  [python]
def exp_map_SE3(tangent_vector: TensorType["b", 6]) -> TensorType["b", 3, 4]:  # pylint: disable=invalid-name
    """Compute the exponential map `se(3) -> SE(3)`.

    This can be used for learning pose deltas on `SE(3)`.

    Args:
        tangent_vector: A tangent vector from `se(3)`.

    Returns:
        [R|t] tranformation matrices.
    """

    tangent_vector_lin = tangent_vector[:, :3].view(-1, 3, 1)
    tangent_vector_ang = tangent_vector[:, 3:].view(-1, 3, 1)

    theta = torch.norm(tangent_vector_ang, dim=1).unsqueeze(1)
    theta2 = theta**2
    theta3 = theta**3

    near_zero = theta < 1e-2
    non_zero = torch.ones(1, dtype=tangent_vector.dtype, device=tangent_vector.device)
    theta_nz = torch.where(near_zero, non_zero, theta)
    theta2_nz = torch.where(near_zero, non_zero, theta2)
    theta3_nz = torch.where(near_zero, non_zero, theta3)

    # Compute the rotation
    sine = theta.sin()
    cosine = torch.where(near_zero, 8 / (4 + theta2) - 1, theta.cos())
    sine_by_theta = torch.where(near_zero, 0.5 * cosine + 0.5, sine / theta_nz)
    one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 * sine_by_theta, (1 - cosine) / theta2_nz)
    ret = torch.zeros(tangent_vector.shape[0], 3, 4).to(dtype=tangent_vector.dtype, device=tangent_vector.device)
    ret[:, :3, :3] = one_minus_cosine_by_theta2 * tangent_vector_ang @ tangent_vector_ang.transpose(1, 2)

    ret[:, 0, 0] += cosine.view(-1)
    ret[:, 1, 1] += cosine.view(-1)
    ret[:, 2, 2] += cosine.view(-1)
    temp = sine_by_theta.view(-1, 1) * tangent_vector_ang.view(-1, 3)
    ret[:, 0, 1] -= temp[:, 2]
    ret[:, 1, 0] += temp[:, 2]
    ret[:, 0, 2] += temp[:, 1]
    ret[:, 2, 0] -= temp[:, 1]
    ret[:, 1, 2] -= temp[:, 0]
    ret[:, 2, 1] += temp[:, 0]

    # Compute the translation
    sine_by_theta = torch.where(near_zero, 1 - theta2 / 6, sine_by_theta)
    one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 - theta2 / 24, one_minus_cosine_by_theta2)
    theta_minus_sine_by_theta3_t = torch.where(near_zero, 1.0 / 6 - theta2 / 120, (theta - sine) / theta3_nz)

    ret[:, :, 3:] = sine_by_theta * tangent_vector_lin
    ret[:, :, 3:] += one_minus_cosine_by_theta2 * torch.cross(tangent_vector_ang, tangent_vector_lin, dim=1)
    ret[:, :, 3:] += theta_minus_sine_by_theta3_t * (
        tangent_vector_ang @ (tangent_vector_ang.transpose(1, 2) @ tangent_vector_lin)
    )
    return ret
//  [python]
def exp_map_SE3(tangent_vector: TensorType["b", 6]) -> TensorType["b", 3, 4]:  # pylint: disable=invalid-name
    """Compute the exponential map `se(3) -> SE(3)`.

    This can be used for learning pose deltas on `SE(3)`.

    Args:
        tangent_vector: A tangent vector from `se(3)`.

    Returns:
        [R|t] tranformation matrices.
    """

    tangent_vector_lin = tangent_vector[:, :3].view(-1, 3, 1)
    tangent_vector_ang = tangent_vector[:, 3:].view(-1, 3, 1)

    theta = torch.linalg.norm(tangent_vector_ang, dim=1).unsqueeze(1)
    theta2 = theta**2
    theta3 = theta**3

    near_zero = theta < 1e-2
    non_zero = torch.ones(1, dtype=tangent_vector.dtype, device=tangent_vector.device)
    theta_nz = torch.where(near_zero, non_zero, theta)
    theta2_nz = torch.where(near_zero, non_zero, theta2)
    theta3_nz = torch.where(near_zero, non_zero, theta3)

    # Compute the rotation
    sine = theta.sin()
    cosine = torch.where(near_zero, 8 / (4 + theta2) - 1, theta.cos())
    sine_by_theta = torch.where(near_zero, 0.5 * cosine + 0.5, sine / theta_nz)
    one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 * sine_by_theta, (1 - cosine) / theta2_nz)
    ret = torch.zeros(tangent_vector.shape[0], 3, 4).to(dtype=tangent_vector.dtype, device=tangent_vector.device)
    ret[:, :3, :3] = one_minus_cosine_by_theta2 * tangent_vector_ang @ tangent_vector_ang.transpose(1, 2)

    ret[:, 0, 0] += cosine.view(-1)
    ret[:, 1, 1] += cosine.view(-1)
    ret[:, 2, 2] += cosine.view(-1)
    temp = sine_by_theta.view(-1, 1) * tangent_vector_ang.view(-1, 3)
    ret[:, 0, 1] -= temp[:, 2]
    ret[:, 1, 0] += temp[:, 2]
    ret[:, 0, 2] += temp[:, 1]
    ret[:, 2, 0] -= temp[:, 1]
    ret[:, 1, 2] -= temp[:, 0]
    ret[:, 2, 1] += temp[:, 0]

    # Compute the translation
    sine_by_theta = torch.where(near_zero, 1 - theta2 / 6, sine_by_theta)
    one_minus_cosine_by_theta2 = torch.where(near_zero, 0.5 - theta2 / 24, one_minus_cosine_by_theta2)
    theta_minus_sine_by_theta3_t = torch.where(near_zero, 1.0 / 6 - theta2 / 120, (theta - sine) / theta3_nz)

    ret[:, :, 3:] = sine_by_theta * tangent_vector_lin
    ret[:, :, 3:] += one_minus_cosine_by_theta2 * torch.cross(tangent_vector_ang, tangent_vector_lin, dim=1)
    ret[:, :, 3:] += theta_minus_sine_by_theta3_t * (
        tangent_vector_ang @ (tangent_vector_ang.transpose(1, 2) @ tangent_vector_lin)
    )
    return ret
//  [python]
def OTExperiments(dim_z=5, m=10, n=10, batch=10, iters=500, trials=10, approx=False):

    learning_curves = [[] for i in range(trials)]
    cosine_sim_curves = [[] for i in range(trials)]

    for trial in range(trials):
        # prepare data and model
        torch.manual_seed(22 + trial)
        P_true = sinkhorn(torch.randn((batch, m, n), dtype=torch.float, requires_grad=False))
        z_init = torch.randn((batch, dim_z), dtype=torch.float, requires_grad=False)

        InstrumentedApproxOptimalTransportFcn.trace = cosine_sim_curves[trial]

        model = OTNetwork(dim_z, m, n, approx=approx)
        optimizer = torch.optim.AdamW(model.parameters(), lr=1.0e-3)

        # do optimisation
        for i in range(iters):
            optimizer.zero_grad(set_to_none=True)
            P_pred = model(z_init)
            loss = torch.norm(P_pred - P_true)
            learning_curves[trial].append(float(loss.item()))
            loss.backward()
            optimizer.step()

            if (i % 100 == 0):
                print("{: 3} {: 6}: {}".format(trial, i, loss.item()))

        InstrumentedApproxOptimalTransportFcn.trace = None

    return learning_curves, cosine_sim_curves
//  [python]
def OTExperiments(dim_z=5, m=10, n=10, batch=10, iters=500, trials=10, approx=False):

    learning_curves = [[] for i in range(trials)]
    cosine_sim_curves = [[] for i in range(trials)]

    for trial in range(trials):
        # prepare data and model
        torch.manual_seed(22 + trial)
        P_true = sinkhorn(torch.randn((batch, m, n), dtype=torch.float, requires_grad=False))
        z_init = torch.randn((batch, dim_z), dtype=torch.float, requires_grad=False)

        InstrumentedApproxOptimalTransportFcn.trace = cosine_sim_curves[trial]

        model = OTNetwork(dim_z, m, n, approx=approx)
        optimizer = torch.optim.AdamW(model.parameters(), lr=1.0e-3)

        # do optimisation
        for i in range(iters):
            optimizer.zero_grad(set_to_none=True)
            P_pred = model(z_init)
            loss = torch.linalg.norm(P_pred - P_true)
            learning_curves[trial].append(float(loss.item()))
            loss.backward()
            optimizer.step()

            if (i % 100 == 0):
                print("{: 3} {: 6}: {}".format(trial, i, loss.item()))

        InstrumentedApproxOptimalTransportFcn.trace = None

    return learning_curves, cosine_sim_curves
//  [python]
def pixel_norm(x, eps=1e-6):
    """Pixel Normalization.

    This normalization is proposed in:
    Progressive Growing of GANs for Improved Quality, Stability, and Variation

    Args:
        x (torch.Tensor): Tensor to be normalized.
        eps (float, optional): Epsilon to avoid dividing zero.
            Defaults to 1e-6.

    Returns:
        torch.Tensor: Normalized tensor.
    """
    if torch.__version__ >= '1.7.0':
        norm = torch.linalg.norm(x, ord=2, dim=1, keepdim=True)
    # support older pytorch version
    else:
        norm = torch.norm(x, p=2, dim=1, keepdim=True)
    norm = norm / torch.sqrt(torch.tensor(x.shape[1]).to(x))

    return x / (norm + eps)
//  [python]
def prepare_ray(self, data_dict):
        def normalize_scene(coord, scene_bbox):
            scene_bbox = torch.tensor(scene_bbox).to(coord)
            norm_coord = (coord - scene_bbox[:3]) / (scene_bbox[3:] - scene_bbox[:3])
            return norm_coord

        condition = data_dict["condition"][0]
        dataset_idx = self.conditions.index(condition)
        scene_bbox = self.scene_bbox[dataset_idx]
        # [0, 1]
        ray_start = normalize_scene(data_dict["ray_start"], scene_bbox)
        ray_end = normalize_scene(data_dict["ray_end"], scene_bbox)
        ray_d = F.normalize(ray_end - ray_start, dim=-1)
        ray_depth = torch.norm(ray_end - ray_start, dim=-1, keepdim=True)

        ray_dict = {
            "ray_offset": data_dict["ray_offset"],
            "ray_o": ray_start,
            "ray_d": ray_d,
            "depth": ray_depth,
        }
        if "ray_color" in data_dict.keys():
            ray_dict["rgb"] = data_dict["ray_color"]
        if "ray_segment" in data_dict.keys() and self.render_semantic:
            assert (
                len(set(data_dict["condition"])) == 1
            ), "assume same condition in one batch"
            condition = data_dict["condition"][0]
            segmen2semantic = self.class_embedding[
                self.valid_index[self.conditions.index(condition)], :
            ]
            ray_segment = data_dict["ray_segment"]
            ray_semantic = segmen2semantic[ray_segment.long()]
            ray_dict["semantic"] = ray_semantic
            ray_dict["segment"] = ray_segment
        return ray_dict
//  [python]
def prepare_ray(self, data_dict):
        def normalize_scene(coord, scene_bbox):
            scene_bbox = torch.tensor(scene_bbox).to(coord)
            norm_coord = (coord - scene_bbox[:3]) / (scene_bbox[3:] - scene_bbox[:3])
            return norm_coord

        condition = data_dict["condition"][0]
        dataset_idx = self.conditions.index(condition)
        scene_bbox = self.scene_bbox[dataset_idx]
        # [0, 1]
        ray_start = normalize_scene(data_dict["ray_start"], scene_bbox)
        ray_end = normalize_scene(data_dict["ray_end"], scene_bbox)
        ray_d = F.normalize(ray_end - ray_start, dim=-1)
        ray_depth = torch.linalg.norm(ray_end - ray_start, dim=-1, keepdim=True)

        ray_dict = {
            "ray_offset": data_dict["ray_offset"],
            "ray_o": ray_start,
            "ray_d": ray_d,
            "depth": ray_depth,
        }
        if "ray_color" in data_dict.keys():
            ray_dict["rgb"] = data_dict["ray_color"]
        if "ray_segment" in data_dict.keys() and self.render_semantic:
            assert (
                len(set(data_dict["condition"])) == 1
            ), "assume same condition in one batch"
            condition = data_dict["condition"][0]
            segmen2semantic = self.class_embedding[
                self.valid_index[self.conditions.index(condition)], :
            ]
            ray_segment = data_dict["ray_segment"]
            ray_semantic = segmen2semantic[ray_segment.long()]
            ray_dict["semantic"] = ray_semantic
            ray_dict["segment"] = ray_segment
        return ray_dict
//  [python]
def cheirality_check(models, points):
    epipoles = torch.cross(models[..., 0, :], models[..., 2, :])
    degenerate_mask = torch.norm(epipoles, dim=-1) < 1e-6
    if torch.any(degenerate_mask):
        epi2 = torch.cross(models[degenerate_mask][..., 1, :], models[degenerate_mask][..., 2, :])
        epipoles[degenerate_mask] = epi2
    s1 = models[..., 0, 0].unsqueeze(1) * points[..., 2] + models[..., 1, 0].unsqueeze(1) * points[..., 3] + models[..., 2, 0].unsqueeze(1)
    s2 = epipoles[..., 1].unsqueeze(1) - epipoles[..., 2].unsqueeze(1) * points[..., 1]
    s = s1 * s2 > 0

    not_ok = s.float().mean(dim=-1) > 0.5
    s[not_ok] = ~s[not_ok]
    return s