jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_linalg_norm(
        self, test_name, input_shape, ord=None, dim=None, keepdims=False
    ):
        class TestModule(torch.nn.Module):
            def __init__(self, ord, dim, keepdims):
                super().__init__()

            def forward(self, x: torch.Tensor) -> torch.Tensor:
                return torch.linalg.norm(x, ord, dim, keepdims)

        model = TestModule(ord, dim, keepdims).cuda().half()
        inputs = [
            torch.randn(input_shape).half().cuda(),
        ]

        self.run_test(model, inputs, expected_ops={acc_ops.linalg_norm})
//  [python]
def pre_compute(self):
        # This method is auxiliary of the `get_solution` method (below)
        # It extracts some computation irrelevant to `area` to reduce computational time
        # You'll need to first understand the theories to solve truss problems and read this method together with the `get_solution` method to understand
        dis_vec = self.coords[self.con[:, 1], :] - self.coords[self.con[:, 0], :]
        self.length = torch.norm(dis_vec, dim=-1)  # length of bars
        self.T = dis_vec / self.length.view(-1, 1)  # distance unit vector or transformation vector
        self.matrix = torch.kron(torch.tensor([[1, -1], [-1, 1]], device=self.device), torch.einsum('ij,ik->ijk', self.T, self.T))
        self.idx = torch.cat((torch.arange(0, self.n_dim, device=self.device) + (self.con[:, 0] * self.n_dim).view(-1, 1),
                              torch.arange(0, self.n_dim, device=self.device) + (self.con[:, 1] * self.n_dim).view(-1, 1)), dim=1)
        self.free_list = self.free_mask.flatten()
//  [python]
def pre_compute(self):
        # This method is auxiliary of the `get_solution` method (below)
        # It extracts some computation irrelevant to `area` to reduce computational time
        # You'll need to first understand the theories to solve truss problems and read this method together with the `get_solution` method to understand
        dis_vec = self.coords[self.con[:, 1], :] - self.coords[self.con[:, 0], :]
        self.length = torch.linalg.norm(dis_vec, dim=-1)  # length of bars
        self.T = dis_vec / self.length.view(-1, 1)  # distance unit vector or transformation vector
        self.matrix = torch.kron(torch.tensor([[1, -1], [-1, 1]], device=self.device), torch.einsum('ij,ik->ijk', self.T, self.T))
        self.idx = torch.cat((torch.arange(0, self.n_dim, device=self.device) + (self.con[:, 0] * self.n_dim).view(-1, 1),
                              torch.arange(0, self.n_dim, device=self.device) + (self.con[:, 1] * self.n_dim).view(-1, 1)), dim=1)
        self.free_list = self.free_mask.flatten()
//  [python]
def with_edge_vectors(data: Type, with_lengths: bool = True) -> Type:
    """Compute the edge displacement vectors for a graph.

    If ``data.pos.requires_grad`` and/or ``data.cell.requires_grad``, this
    method will return edge vectors correctly connected in the autograd graph.

    Returns:
        Tensor [n_edges, 3] edge displacement vectors
    """
    if _keys.EDGE_VECTORS_KEY in data:
        if with_lengths and _keys.EDGE_LENGTH_KEY not in data:
            data[_keys.EDGE_LENGTH_KEY] = torch.norm(
                data[_keys.EDGE_VECTORS_KEY], dim=-1
            )
        return data
    else:
        # Build it dynamically
        # Note that this is
        # (1) backwardable, because everything (pos, cell, shifts)
        #     is Tensors.
        # (2) works on a Batch constructed from AtomicData
        pos = data[_keys.POSITIONS_KEY]
        edge_index = data[_keys.EDGE_INDEX_KEY]
        edge_vec = pos[edge_index[1]] - pos[edge_index[0]]
        if _keys.CELL_KEY in data:
            # ^ note that to save time we don't check that the edge_cell_shifts are trivial if no cell is provided; we just assume they are either not present or all zero.
            # -1 gives a batch dim no matter what
            cell = data[_keys.CELL_KEY].view(-1, 3, 3)
            edge_cell_shift = data[_keys.EDGE_CELL_SHIFT_KEY]
            if cell.shape[0] > 1:
                batch = data[_keys.BATCH_KEY]
                # Cell has a batch dimension
                # note the ASE cell vectors as rows convention
                edge_vec = edge_vec + torch.einsum(
                    "ni,nij->nj", edge_cell_shift, cell[batch[edge_index[0]]]
                )
                # TODO: is there a more efficient way to do the above without
                # creating an [n_edge] and [n_edge, 3, 3] tensor?
            else:
                # Cell has either no batch dimension, or a useless one,
                # so we can avoid creating the large intermediate cell tensor.
                # Note that we do NOT check that the batch array, if it is present,
                # is trivial — but this does need to be consistent.
                edge_vec = edge_vec + torch.einsum(
                    "ni,ij->nj",
                    edge_cell_shift,
                    cell.squeeze(0),  # remove batch dimension
                )
        data[_keys.EDGE_VECTORS_KEY] = edge_vec
        if with_lengths:
            data[_keys.EDGE_LENGTH_KEY] = torch.norm(edge_vec, dim=-1)
        return data
//  [python]
def with_edge_vectors(data: Type, with_lengths: bool = True) -> Type:
    """Compute the edge displacement vectors for a graph.

    If ``data.pos.requires_grad`` and/or ``data.cell.requires_grad``, this
    method will return edge vectors correctly connected in the autograd graph.

    Returns:
        Tensor [n_edges, 3] edge displacement vectors
    """
    if _keys.EDGE_VECTORS_KEY in data:
        if with_lengths and _keys.EDGE_LENGTH_KEY not in data:
            data[_keys.EDGE_LENGTH_KEY] = torch.linalg.norm(
                data[_keys.EDGE_VECTORS_KEY], dim=-1
            )
        return data
    else:
        # Build it dynamically
        # Note that this is
        # (1) backwardable, because everything (pos, cell, shifts)
        #     is Tensors.
        # (2) works on a Batch constructed from AtomicData
        pos = data[_keys.POSITIONS_KEY]
        edge_index = data[_keys.EDGE_INDEX_KEY]
        edge_vec = pos[edge_index[1]] - pos[edge_index[0]]
        if _keys.CELL_KEY in data:
            # ^ note that to save time we don't check that the edge_cell_shifts are trivial if no cell is provided; we just assume they are either not present or all zero.
            # -1 gives a batch dim no matter what
            cell = data[_keys.CELL_KEY].view(-1, 3, 3)
            edge_cell_shift = data[_keys.EDGE_CELL_SHIFT_KEY]
            if cell.shape[0] > 1:
                batch = data[_keys.BATCH_KEY]
                # Cell has a batch dimension
                # note the ASE cell vectors as rows convention
                edge_vec = edge_vec + torch.einsum(
                    "ni,nij->nj", edge_cell_shift, cell[batch[edge_index[0]]]
                )
                # TODO: is there a more efficient way to do the above without
                # creating an [n_edge] and [n_edge, 3, 3] tensor?
            else:
                # Cell has either no batch dimension, or a useless one,
                # so we can avoid creating the large intermediate cell tensor.
                # Note that we do NOT check that the batch array, if it is present,
                # is trivial — but this does need to be consistent.
                edge_vec = edge_vec + torch.einsum(
                    "ni,ij->nj",
                    edge_cell_shift,
                    cell.squeeze(0),  # remove batch dimension
                )
        data[_keys.EDGE_VECTORS_KEY] = edge_vec
        if with_lengths:
            data[_keys.EDGE_LENGTH_KEY] = torch.linalg.norm(edge_vec, dim=-1)
        return data
//  [python]
def create_table_neighbour_code_to_surface_area(spacing_mm, device=None):
    """
    Returns an array mapping neighbourhood code to the surface elements area.
    Adapted from https://github.com/deepmind/surface-distance

    Note that the normals encode the initial surface area. This function computes
    the area corresponding to the given `spacing`.

    Args:
        spacing_mm: a sequence of 3 numbers. Voxel spacing along the first 3 spatial axes.
        device: device to put the table on.

    Returns:
        An array of size 256, mapping neighbourhood code to the surface area.
        ENCODING_KERNEL[3] which is the kernel used to compute the neighbourhood code.
    """
    spacing_mm = ensure_tuple_rep(spacing_mm, 3)
    # compute the area for all 256 possible surface elements given a 2x2x2 neighbourhood according to the spacing_mm
    c = _get_neighbour_code_to_normals_table(device)
    s = torch.as_tensor(
        [[[spacing_mm[1] * spacing_mm[2], spacing_mm[0] * spacing_mm[2], spacing_mm[0] * spacing_mm[1]]]],
        device=device,
        dtype=c.dtype,
    )
    norm = torch.norm(c * s, dim=-1)
    neighbour_code_to_surface_area = norm.sum(-1)
    return neighbour_code_to_surface_area, torch.as_tensor([[ENCODING_KERNEL[3]]], device=device)
//  [python]
def create_table_neighbour_code_to_surface_area(spacing_mm, device=None):
    """
    Returns an array mapping neighbourhood code to the surface elements area.
    Adapted from https://github.com/deepmind/surface-distance

    Note that the normals encode the initial surface area. This function computes
    the area corresponding to the given `spacing`.

    Args:
        spacing_mm: a sequence of 3 numbers. Voxel spacing along the first 3 spatial axes.
        device: device to put the table on.

    Returns:
        An array of size 256, mapping neighbourhood code to the surface area.
        ENCODING_KERNEL[3] which is the kernel used to compute the neighbourhood code.
    """
    spacing_mm = ensure_tuple_rep(spacing_mm, 3)
    # compute the area for all 256 possible surface elements given a 2x2x2 neighbourhood according to the spacing_mm
    c = _get_neighbour_code_to_normals_table(device)
    s = torch.as_tensor(
        [[[spacing_mm[1] * spacing_mm[2], spacing_mm[0] * spacing_mm[2], spacing_mm[0] * spacing_mm[1]]]],
        device=device,
        dtype=c.dtype,
    )
    norm = torch.linalg.norm(c * s, dim=-1)
    neighbour_code_to_surface_area = norm.sum(-1)
    return neighbour_code_to_surface_area, torch.as_tensor([[ENCODING_KERNEL[3]]], device=device)
//  [python]
def test_gaussian_penalty(self):
        for dtype in (torch.float, torch.double):
            tkwargs = {"device": self.device, "dtype": dtype}
            init_point = torch.tensor([1.0, 1.0, 1.0], **tkwargs)
            l2_module = L2Penalty(init_point=init_point)

            # testing a batch of two points
            sample_point = torch.tensor([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]], **tkwargs)

            diff_norm_squared = torch.norm((sample_point - init_point), p=2, dim=-1) ** 2
            real_value = diff_norm_squared.max(dim=-1).values
            computed_value = l2_module(sample_point)
            self.assertEqual(computed_value.item(), real_value.item())
//  [python]
def test_gaussian_penalty(self):
        for dtype in (torch.float, torch.double):
            tkwargs = {"device": self.device, "dtype": dtype}
            init_point = torch.tensor([1.0, 1.0, 1.0], **tkwargs)
            l2_module = L2Penalty(init_point=init_point)

            # testing a batch of two points
            sample_point = torch.tensor([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]], **tkwargs)

            diff_norm_squared = (
                torch.linalg.norm((sample_point - init_point), ord=2, dim=-1) ** 2
            )
            real_value = diff_norm_squared.max(dim=-1).values
            computed_value = l2_module(sample_point)
            self.assertEqual(computed_value.item(), real_value.item())
//  [python]
def test_l1_penalty(self):
        for dtype in (torch.float, torch.double):
            init_point = torch.tensor([1.0, 1.0, 1.0], device=self.device, dtype=dtype)
            l1_module = L1Penalty(init_point=init_point)

            # testing a batch of two points
            sample_point = torch.tensor(
                [[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]], device=self.device, dtype=dtype
            )

            diff_l1_norm = torch.norm((sample_point - init_point), p=1, dim=-1)
            real_value = diff_l1_norm.max(dim=-1).values
            computed_value = l1_module(sample_point)
            self.assertEqual(computed_value.item(), real_value.item())
//  [python]
def test_l1_penalty(self):
        for dtype in (torch.float, torch.double):
            init_point = torch.tensor([1.0, 1.0, 1.0], device=self.device, dtype=dtype)
            l1_module = L1Penalty(init_point=init_point)

            # testing a batch of two points
            sample_point = torch.tensor(
                [[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]], device=self.device, dtype=dtype
            )

            diff_l1_norm = torch.linalg.norm((sample_point - init_point), ord=1, dim=-1)
            real_value = diff_l1_norm.max(dim=-1).values
            computed_value = l1_module(sample_point)
            self.assertEqual(computed_value.item(), real_value.item())
//  [python]
def test_gaussian_penalty(self):
        for dtype in (torch.float, torch.double):
            tkwargs = {"device": self.device, "dtype": dtype}
            init_point = torch.tensor([1.0, 1.0, 1.0], **tkwargs)
            sigma = 0.1
            gaussian_module = GaussianPenalty(init_point=init_point, sigma=sigma)

            # testing a batch of two points
            sample_point = torch.tensor([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]], **tkwargs)

            diff_norm_squared = (
                torch.norm((sample_point - init_point), p=2, dim=-1) ** 2
            )
            max_l2_distance = diff_norm_squared.max(dim=-1).values
            real_value = torch.exp(max_l2_distance / 2 / sigma**2)
            computed_value = gaussian_module(sample_point)
            self.assertEqual(computed_value.item(), real_value.item())
//  [python]
def test_gaussian_penalty(self):
        for dtype in (torch.float, torch.double):
            tkwargs = {"device": self.device, "dtype": dtype}
            init_point = torch.tensor([1.0, 1.0, 1.0], **tkwargs)
            sigma = 0.1
            gaussian_module = GaussianPenalty(init_point=init_point, sigma=sigma)

            # testing a batch of two points
            sample_point = torch.tensor([[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]], **tkwargs)

            diff_norm_squared = (
                torch.linalg.norm((sample_point - init_point), ord=2, dim=-1) ** 2
            )
            max_l2_distance = diff_norm_squared.max(dim=-1).values
            real_value = torch.exp(max_l2_distance / 2 / sigma**2)
            computed_value = gaussian_module(sample_point)
            self.assertEqual(computed_value.item(), real_value.item())
//  [python]
def axis_angle_from_quat(quat: torch.Tensor, eps: float = 1.0e-6) -> torch.Tensor:
    """Convert rotations given as quaternions to axis/angle.

    Args:
        quat: The quaternion orientation in (w, x, y, z). Shape is (..., 4).
        eps: The tolerance for Taylor approximation. Defaults to 1.0e-6.

    Returns:
        Rotations given as a vector in axis angle form. Shape is (..., 3).
        The vector's magnitude is the angle turned anti-clockwise in radians around the vector's direction.

    Reference:
        https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/transforms/rotation_conversions.py#L526-L554
    """
    # Modified to take in quat as [q_w, q_x, q_y, q_z]
    # Quaternion is [q_w, q_x, q_y, q_z] = [cos(theta/2), n_x * sin(theta/2), n_y * sin(theta/2), n_z * sin(theta/2)]
    # Axis-angle is [a_x, a_y, a_z] = [theta * n_x, theta * n_y, theta * n_z]
    # Thus, axis-angle is [q_x, q_y, q_z] / (sin(theta/2) / theta)
    # When theta = 0, (sin(theta/2) / theta) is undefined
    # However, as theta --> 0, we can use the Taylor approximation 1/2 - theta^2 / 48
    quat = quat * (1.0 - 2.0 * (quat[..., 0:1] < 0.0))
    mag = torch.norm(quat[..., 1:], dim=-1)
    half_angle = torch.atan2(mag, quat[..., 0])
    angle = 2.0 * half_angle
    # check whether to apply Taylor approximation
    sin_half_angles_over_angles = torch.where(
        angle.abs() > eps, torch.sin(half_angle) / angle, 0.5 - angle * angle / 48
    )
    return quat[..., 1:4] / sin_half_angles_over_angles.unsqueeze(-1)
//  [python]
def axis_angle_from_quat(quat: torch.Tensor, eps: float = 1.0e-6) -> torch.Tensor:
    """Convert rotations given as quaternions to axis/angle.

    Args:
        quat: The quaternion orientation in (w, x, y, z). Shape is (..., 4).
        eps: The tolerance for Taylor approximation. Defaults to 1.0e-6.

    Returns:
        Rotations given as a vector in axis angle form. Shape is (..., 3).
        The vector's magnitude is the angle turned anti-clockwise in radians around the vector's direction.

    Reference:
        https://github.com/facebookresearch/pytorch3d/blob/main/pytorch3d/transforms/rotation_conversions.py#L526-L554
    """
    # Modified to take in quat as [q_w, q_x, q_y, q_z]
    # Quaternion is [q_w, q_x, q_y, q_z] = [cos(theta/2), n_x * sin(theta/2), n_y * sin(theta/2), n_z * sin(theta/2)]
    # Axis-angle is [a_x, a_y, a_z] = [theta * n_x, theta * n_y, theta * n_z]
    # Thus, axis-angle is [q_x, q_y, q_z] / (sin(theta/2) / theta)
    # When theta = 0, (sin(theta/2) / theta) is undefined
    # However, as theta --> 0, we can use the Taylor approximation 1/2 - theta^2 / 48
    quat = quat * (1.0 - 2.0 * (quat[..., 0:1] < 0.0))
    mag = torch.linalg.norm(quat[..., 1:], dim=-1)
    half_angle = torch.atan2(mag, quat[..., 0])
    angle = 2.0 * half_angle
    # check whether to apply Taylor approximation
    sin_half_angles_over_angles = torch.where(
        angle.abs() > eps, torch.sin(half_angle) / angle, 0.5 - angle * angle / 48
    )
    return quat[..., 1:4] / sin_half_angles_over_angles.unsqueeze(-1)
//  [python]
def test_l1_penalty(self):
        for dtype in (torch.float, torch.double):
            init_point = torch.tensor([1.0, 1.0, 1.0], device=self.device, dtype=dtype)
            l1_module = L1PenaltyObjective(init_point=init_point)

            # testing a batch of two points
            sample_point = torch.tensor(
                [[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]], device=self.device, dtype=dtype
            )

            real_values = torch.norm(
                (sample_point - init_point), p=1, dim=-1
            ).unsqueeze(dim=0)
            computed_values = l1_module(sample_point)
            self.assertTrue(torch.equal(real_values, computed_values))
//  [python]
def test_l1_penalty(self):
        for dtype in (torch.float, torch.double):
            init_point = torch.tensor([1.0, 1.0, 1.0], device=self.device, dtype=dtype)
            l1_module = L1PenaltyObjective(init_point=init_point)

            # testing a batch of two points
            sample_point = torch.tensor(
                [[1.0, 2.0, 3.0], [2.0, 3.0, 4.0]], device=self.device, dtype=dtype
            )

            real_values = torch.linalg.norm(
                (sample_point - init_point), ord=1, dim=-1
            ).unsqueeze(dim=0)
            computed_values = l1_module(sample_point)
            self.assertTrue(torch.equal(real_values, computed_values))
//  [python]
def exp_map_SE3(tangent_vector: Float[Tensor, "b 6"]) -> Float[Tensor, "b 3 4"]:
    """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] transformation 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: Float[Tensor, "b 6"]) -> Float[Tensor, "b 3 4"]:
    """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] transformation 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 verify_speakers(self, path2audio_file1, path2audio_file2, threshold=0.7):
        """
        Verify if two audio files are from the same speaker or not.

        Args:
            path2audio_file1: path to audio wav file of speaker 1
            path2audio_file2: path to audio wav file of speaker 2
            threshold: cosine similarity score used as a threshold to distinguish two embeddings (default = 0.7)

        Returns:
            True if both audio files are from same speaker, False otherwise
        """
        embs1 = self.get_embedding(path2audio_file1).squeeze()
        embs2 = self.get_embedding(path2audio_file2).squeeze()
        # Length Normalize
        X = embs1 / torch.norm(embs1)
        Y = embs2 / torch.norm(embs2)
        # Score
        similarity_score = torch.dot(X, Y) / ((torch.dot(X, X) * torch.dot(Y, Y)) ** 0.5)
        similarity_score = (similarity_score + 1) / 2
        # Decision
        if similarity_score >= threshold:
            logging.info(" two audio files are from same speaker")
            return True
        else:
            logging.info(" two audio files are from different speakers")
            return False
//  [python]
def verify_speakers(self, path2audio_file1, path2audio_file2, threshold=0.7):
        """
        Verify if two audio files are from the same speaker or not.

        Args:
            path2audio_file1: path to audio wav file of speaker 1
            path2audio_file2: path to audio wav file of speaker 2
            threshold: cosine similarity score used as a threshold to distinguish two embeddings (default = 0.7)

        Returns:
            True if both audio files are from same speaker, False otherwise
        """
        embs1 = self.get_embedding(path2audio_file1).squeeze()
        embs2 = self.get_embedding(path2audio_file2).squeeze()
        # Length Normalize
        X = embs1 / torch.linalg.norm(embs1)
        Y = embs2 / torch.linalg.norm(embs2)
        # Score
        similarity_score = torch.dot(X, Y) / ((torch.dot(X, X) * torch.dot(Y, Y)) ** 0.5)
        similarity_score = (similarity_score + 1) / 2
        # Decision
        if similarity_score >= threshold:
            logging.info(" two audio files are from same speaker")
            return True
        else:
            logging.info(" two audio files are from different speakers")
            return False
//  [python]
def render_core_outside(self, rays_o, rays_d, z_vals, sample_dist, nerf, background_rgb=None):
        batch_size, n_samples = z_vals.shape

        dists = z_vals[..., 1:] - z_vals[..., :-1]
        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)

        mid_z_vals = z_vals + dists * 0.5

        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # batch_size, n_samples, 3

        if self.n_outside > 0:
            dis_to_center = torch.norm(pts, dim=-1, keepdim=True).clip(1.0, 1e10)
            pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center], dim=-1)  # batch_size, n_samples, 4

        dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)

        pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
        dirs = dirs.reshape(-1, 3)

        raw, sampled_color = nerf(pts, dirs)
        alpha = 1.0 - torch.exp(-F.relu(raw.reshape(batch_size, n_samples)) * dists)
        alpha = alpha.reshape(batch_size, n_samples)
        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:,
                          :-1]  # n_rays, n_samples
        sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
        color = (weights[:, :, None] * sampled_color).sum(dim=1)
        if background_rgb is not None:
            color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))

        return {
            'color': color,
            'sampled_color': sampled_color,
            'alpha': alpha,
            'weights': weights,
        }
//  [python]
def render_core_outside(self, rays_o, rays_d, z_vals, sample_dist, nerf, background_rgb=None):
        batch_size, n_samples = z_vals.shape

        dists = z_vals[..., 1:] - z_vals[..., :-1]
        dists = torch.cat([dists, torch.Tensor([sample_dist]).expand(dists[..., :1].shape)], -1)

        mid_z_vals = z_vals + dists * 0.5

        pts = rays_o[:, None, :] + rays_d[:, None, :] * mid_z_vals[..., :, None]  # batch_size, n_samples, 3

        if self.n_outside > 0:
            dis_to_center = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=True).clip(1.0, 1e10)
            pts = torch.cat([pts / dis_to_center, 1.0 / dis_to_center], dim=-1)  # batch_size, n_samples, 4

        dirs = rays_d[:, None, :].expand(batch_size, n_samples, 3)

        pts = pts.reshape(-1, 3 + int(self.n_outside > 0))
        dirs = dirs.reshape(-1, 3)

        raw, sampled_color = nerf(pts, dirs)
        alpha = 1.0 - torch.exp(-F.relu(raw.reshape(batch_size, n_samples)) * dists)
        alpha = alpha.reshape(batch_size, n_samples)
        weights = alpha * torch.cumprod(torch.cat([torch.ones([batch_size, 1]), 1. - alpha + 1e-7], -1), -1)[:,
                          :-1]  # n_rays, n_samples
        sampled_color = sampled_color.reshape(batch_size, n_samples, 3)
        color = (weights[:, :, None] * sampled_color).sum(dim=1)
        if background_rgb is not None:
            color = color + background_rgb * (1.0 - weights.sum(dim=-1, keepdim=True))

        return {
            'color': color,
            'sampled_color': sampled_color,
            'alpha': alpha,
            'weights': weights,
        }
//  [python]
def up_sample_no_occ_aware(self, rays_o, rays_d, z_vals, udf, sample_dist,
                               n_importance, inv_s, beta, gamma, ):
        """
        Different with NeuS, here we sample more points at all possible surfaces where udf is close to 0;
        Since unlike that SDF has clear sign changes, UDF sampling may miss the true surfaces
        """
        device = z_vals.device
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)

        udf = udf.reshape(batch_size, n_samples)

        dists = z_vals[..., 1:] - z_vals[..., :-1]
        dists = torch.cat([dists, torch.Tensor([sample_dist]).to(device).expand(dists[..., :1].shape)], -1)

        dirs = rays_d[:, None, :].expand(pts.shape)

        # * the probability of occlusion
        raw_occ = self.udf2logistic(udf, beta, gamma, 1.0)
        # near 0levelset alpha_acc -> 1 others -> 0
        alpha_occ = 1.0 - torch.exp(-F.relu(raw_occ.reshape(batch_size, n_samples)) * dists)

        z_samples = sample_pdf(z_vals, alpha_occ[:, :-1], n_importance, det=True).detach()

        flag = torch.any(torch.isnan(z_samples)).cpu().numpy().item()
        if flag:
            print("z_vals", z_samples[torch.isnan(z_samples)])
            print('z_vals have nan values')
            pdb.set_trace()

        return z_samples
//  [python]
def up_sample_no_occ_aware(self, rays_o, rays_d, z_vals, udf, sample_dist,
                               n_importance, inv_s, beta, gamma, ):
        """
        Different with NeuS, here we sample more points at all possible surfaces where udf is close to 0;
        Since unlike that SDF has clear sign changes, UDF sampling may miss the true surfaces
        """
        device = z_vals.device
        batch_size, n_samples = z_vals.shape
        pts = rays_o[:, None, :] + rays_d[:, None, :] * z_vals[..., :, None]  # n_rays, n_samples, 3
        radius = torch.linalg.norm(pts, ord=2, dim=-1, keepdim=False)
        inside_sphere = (radius[:, :-1] < 1.0) | (radius[:, 1:] < 1.0)

        udf = udf.reshape(batch_size, n_samples)

        dists = z_vals[..., 1:] - z_vals[..., :-1]
        dists = torch.cat([dists, torch.Tensor([sample_dist]).to(device).expand(dists[..., :1].shape)], -1)

        dirs = rays_d[:, None, :].expand(pts.shape)

        # * the probability of occlusion
        raw_occ = self.udf2logistic(udf, beta, gamma, 1.0)
        # near 0levelset alpha_acc -> 1 others -> 0
        alpha_occ = 1.0 - torch.exp(-F.relu(raw_occ.reshape(batch_size, n_samples)) * dists)

        z_samples = sample_pdf(z_vals, alpha_occ[:, :-1], n_importance, det=True).detach()

        flag = torch.any(torch.isnan(z_samples)).cpu().numpy().item()
        if flag:
            print("z_vals", z_samples[torch.isnan(z_samples)])
            print('z_vals have nan values')
            pdb.set_trace()

        return z_samples
//  [python]
def test_cholesky_ex(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        def run_test(n, batch):
            A = random_hermitian_pd_matrix(n, *batch, dtype=dtype, device=device)
            expected_L = np.linalg.cholesky(A.cpu().numpy())
            expected_info = torch.zeros(A.shape[:-2], dtype=torch.int32, device=device)
            actual_L, actual_info = torch.linalg.cholesky_ex(A)

            # For fp32 individual entries in matrices can differ between PyTorch and NumPy
            # Let's compare the norms of matrices instead
            if A.numel() > 0 and dtype in [torch.float32, torch.complex64]:
                # axis is specified to calculate matrix norm for batched input
                expected_norm = np.linalg.norm(expected_L, ord=1, axis=(-2, -1))
                actual_norm = torch.norm(actual_L, p=1, dim=(-2, -1))
                # Compare the norms with standard tolerances
                self.assertEqual(actual_norm, expected_norm)
                # and individual values with a higher tolerance
                self.assertEqual(actual_L, expected_L, atol=1e-2, rtol=1e-5)
            else:
                self.assertEqual(actual_L, expected_L)
            self.assertEqual(actual_info, expected_info)

        ns = (0, 3, 5)
        batches = ((), (2, ), (2, 1))
        for n, batch in itertools.product(ns, batches):
            run_test(n, batch)
//  [python]
def test_cholesky_ex(self, device, dtype):
        from torch.testing._internal.common_utils import random_hermitian_pd_matrix

        def run_test(n, batch):
            A = random_hermitian_pd_matrix(n, *batch, dtype=dtype, device=device)
            expected_L = np.linalg.cholesky(A.cpu().numpy())
            expected_info = torch.zeros(A.shape[:-2], dtype=torch.int32, device=device)
            actual_L, actual_info = torch.linalg.cholesky_ex(A)

            # For fp32 individual entries in matrices can differ between PyTorch and NumPy
            # Let's compare the norms of matrices instead
            if A.numel() > 0 and dtype in [torch.float32, torch.complex64]:
                # axis is specified to calculate matrix norm for batched input
                expected_norm = np.linalg.norm(expected_L, ord=1, axis=(-2, -1))
                actual_norm = torch.linalg.norm(actual_L, ord=1, axis=(-2, -1))
                # Compare the norms with standard tolerances
                self.assertEqual(actual_norm, expected_norm)
                # and individual values with a higher tolerance
                self.assertEqual(actual_L, expected_L, atol=1e-2, rtol=1e-5)
            else:
                self.assertEqual(actual_L, expected_L)
            self.assertEqual(actual_info, expected_info)

        ns = (0, 3, 5)
        batches = ((), (2, ), (2, 1))
        for n, batch in itertools.product(ns, batches):
            run_test(n, batch)
//  [python]
def linear_distance(self, other_pose: Pose):
        p_distance = torch.norm(self.position - other_pose.position, dim=-1)
        return p_distance
//  [python]
def linear_distance(self, other_pose: Pose):
        p_distance = torch.linalg.norm(self.position - other_pose.position, dim=-1)
        return p_distance
//  [python]
def test_matrix_norm(self, device, dtype):
        # Test only inputs for which torch.linalg.matrix_norm diverges from torch.linalg.norm
        A = make_tensor((2, 2, 2), device, dtype)

        with self.assertRaisesRegex(RuntimeError, r'linalg.matrix_norm\(\):.*must be a matrix.*'):
            torch.linalg.matrix_norm(make_tensor((2,), device, dtype))
        with self.assertRaisesRegex(RuntimeError, r'linalg.matrix_norm\(\):.*must be a 2-tuple.*'):
            torch.linalg.matrix_norm(A, dim=(0,))
        with self.assertRaisesRegex(RuntimeError, r'.*not supported.*'):
            torch.linalg.matrix_norm(A, ord=0)
        with self.assertRaisesRegex(RuntimeError, r'.*not supported.*'):
            torch.linalg.matrix_norm(A, ord=3.0)

        # Test dim=None behavior
        ref = t.norm(A, dim=(-2, -1))
        res = torch.linalg.matrix_norm(A)
        self.assertEqual(ref, res)
//  [python]
def test_matrix_norm(self, device, dtype):
        # Test only inputs for which torch.linalg.matrix_norm diverges from torch.linalg.norm
        A = make_tensor((2, 2, 2), device, dtype)

        with self.assertRaisesRegex(RuntimeError, r'linalg.matrix_norm\(\):.*must be a matrix.*'):
            torch.linalg.matrix_norm(make_tensor((2,), device, dtype))
        with self.assertRaisesRegex(RuntimeError, r'linalg.matrix_norm\(\):.*must be a 2-tuple.*'):
            torch.linalg.matrix_norm(A, dim=(0,))
        with self.assertRaisesRegex(RuntimeError, r'.*not supported.*'):
            torch.linalg.matrix_norm(A, ord=0)
        with self.assertRaisesRegex(RuntimeError, r'.*not supported.*'):
            torch.linalg.matrix_norm(A, ord=3.0)

        # Test dim=None behavior
        ref = torch.linalg.norm(A, dim=(-2, -1))
        res = torch.linalg.matrix_norm(A)
        self.assertEqual(ref, res)
//  [python]
def test_norm_matrix(self, device, dtype):
        def run_test_case(input, ord, dim, keepdim):
            msg = f'input.size()={input.size()}, ord={ord}, dim={dim}, keepdim={keepdim}, dtype={dtype}'
            result = torch.norm(input, ord, dim, keepdim)
            input_numpy = input.cpu().numpy()
            result_numpy = np.linalg.norm(input_numpy, ord, dim, keepdim)

            def check(op):
                result = op(input, ord, dim, keepdim)
                self.assertEqual(result, result_numpy, msg=msg)
                result_out = torch.empty_like(result)
                op(input, ord, dim, keepdim, out=result_out)
                self.assertEqual(result, result_out, msg=msg)

            check(torch.norm)
            if ord is not None and dim is not None:
                check(torch.linalg.matrix_norm)

        ord_matrix = [1, -1, 2, -2, inf, -inf, 'nuc', 'fro']
        S = 10
        test_cases = [
            # input size, p settings, dim
            ((S, S), ord_matrix, None),
            ((S, S), ord_matrix, (0, 1)),
            ((S, S), ord_matrix, (1, 0)),
            ((S, S, S, S), ord_matrix, (2, 0)),
            ((S, S, S, S), ord_matrix, (-1, -2)),
            ((S, S, S, S), ord_matrix, (-1, -3)),
            ((S, S, S, S), ord_matrix, (-3, 2)),
        ]
        L = 1_000

        if dtype == torch.double:
            test_cases.append(((L, L), ord_matrix, None))

        for keepdim in [True, False]:
            for input_size, ord_settings, dim in test_cases:
                input = torch.randn(*input_size, dtype=dtype, device=device)
                for ord in ord_settings:
                    run_test_case(input, ord, dim, keepdim)
//  [python]
def test_norm_matrix(self, device, dtype):
        def run_test_case(input, ord, dim, keepdim):
            msg = f'input.size()={input.size()}, ord={ord}, dim={dim}, keepdim={keepdim}, dtype={dtype}'
            result = torch.linalg.norm(input, ord, dim, keepdim)
            input_numpy = input.cpu().numpy()
            result_numpy = np.linalg.norm(input_numpy, ord, dim, keepdim)

            def check(op):
                result = op(input, ord, dim, keepdim)
                self.assertEqual(result, result_numpy, msg=msg)
                result_out = torch.empty_like(result)
                op(input, ord, dim, keepdim, out=result_out)
                self.assertEqual(result, result_out, msg=msg)

            check(torch.linalg.norm)
            if ord is not None and dim is not None:
                check(torch.linalg.matrix_norm)

        ord_matrix = [1, -1, 2, -2, inf, -inf, 'nuc', 'fro']
        S = 10
        test_cases = [
            # input size, p settings, dim
            ((S, S), ord_matrix, None),
            ((S, S), ord_matrix, (0, 1)),
            ((S, S), ord_matrix, (1, 0)),
            ((S, S, S, S), ord_matrix, (2, 0)),
            ((S, S, S, S), ord_matrix, (-1, -2)),
            ((S, S, S, S), ord_matrix, (-1, -3)),
            ((S, S, S, S), ord_matrix, (-3, 2)),
        ]
        L = 1_000

        if dtype == torch.double:
            test_cases.append(((L, L), ord_matrix, None))

        for keepdim in [True, False]:
            for input_size, ord_settings, dim in test_cases:
                input = torch.randn(*input_size, dtype=dtype, device=device)
                for ord in ord_settings:
                    run_test_case(input, ord, dim, keepdim)
//  [python]
def test_norm_vector_degenerate_shapes(self, device, dtype):
        def run_test_case(input, ord, dim, keepdim):
            msg = f'input.size()={input.size()}, ord={ord}, dim={dim}, keepdim={keepdim}, dtype={dtype}'
            should_error = False
            if ord is not None and ord < 0:
                should_error = True
            elif ord == inf:
                if dim is None or input.size(dim) == 0:
                    should_error = True

            if should_error:
                with self.assertRaises(RuntimeError):
                    torch.norm(input, ord, dim, keepdim)
            else:
                input_numpy = input.cpu().numpy()
                result_numpy = np.linalg.norm(input_numpy, ord, dim, keepdim)
                result = torch.norm(input, ord, dim, keepdim)
                self.assertEqual(result, result_numpy, msg=msg)

        ord_vector = [0, 0.5, 1, 2, 3, inf, -0.5, -1, -2, -3, -inf, None]
        S = 10
        test_cases = [
            # input size, dim
            ((0, ), None),
            ((0, S), 0),
            ((0, S), 1),
            ((S, 0), 0),
            ((S, 0), 1),
        ]
        for keepdim in [True, False]:
            for input_size, dim in test_cases:
                input = torch.randn(*input_size, dtype=dtype, device=device)
                for ord in ord_vector:
                    run_test_case(input, ord, dim, keepdim)
//  [python]
def test_norm_vector_degenerate_shapes(self, device, dtype):
        def run_test_case(input, ord, dim, keepdim):
            msg = f'input.size()={input.size()}, ord={ord}, dim={dim}, keepdim={keepdim}, dtype={dtype}'
            should_error = False
            if ord is not None and ord < 0:
                should_error = True
            elif ord == inf:
                if dim is None or input.size(dim) == 0:
                    should_error = True

            if should_error:
                with self.assertRaises(RuntimeError):
                    torch.linalg.norm(input, ord, dim, keepdim)
            else:
                input_numpy = input.cpu().numpy()
                result_numpy = np.linalg.norm(input_numpy, ord, dim, keepdim)
                result = torch.linalg.norm(input, ord, dim, keepdim)
                self.assertEqual(result, result_numpy, msg=msg)

        ord_vector = [0, 0.5, 1, 2, 3, inf, -0.5, -1, -2, -3, -inf, None]
        S = 10
        test_cases = [
            # input size, dim
            ((0, ), None),
            ((0, S), 0),
            ((0, S), 1),
            ((S, 0), 0),
            ((S, 0), 1),
        ]
        for keepdim in [True, False]:
            for input_size, dim in test_cases:
                input = torch.randn(*input_size, dtype=dtype, device=device)
                for ord in ord_vector:
                    run_test_case(input, ord, dim, keepdim)
//  [python]
def loss_func(x, y, p=2.0):
    yv = y.reshape(x.size()[0], -1)
    xv = x.reshape(x.size()[0], -1)
    diff_norms = torch.norm(xv - yv, p=p, dim=1)
    y_norms = torch.norm(yv, p=p, dim=1)

    return torch.mean(diff_norms / y_norms)
//  [python]
def loss_func(x, y, p=2.0):
    yv = y.reshape(x.size()[0], -1)
    xv = x.reshape(x.size()[0], -1)
    diff_norms = torch.linalg.norm(xv - yv, ord=p, dim=1)
    y_norms = torch.linalg.norm(yv, ord=p, dim=1)

    return torch.mean(diff_norms / y_norms)
//  [python]
def normalize_quaternion(in_quaternion: torch.Tensor) -> torch.Tensor:
    k = torch.sign(in_quaternion[..., 0:1])
    # NOTE: torch sign returns 0 as sign value when value is 0.0
    k = torch.where(k == 0, 1.0, k)
    k2 = k / torch.norm(in_quaternion, dim=-1, keepdim=True)
    # normalize quaternion
    in_q = k2 * in_quaternion
    return in_q
//  [python]
def normalize_quaternion(in_quaternion: torch.Tensor) -> torch.Tensor:
    k = torch.sign(in_quaternion[..., 0:1])
    # NOTE: torch sign returns 0 as sign value when value is 0.0
    k = torch.where(k == 0, 1.0, k)
    k2 = k / torch.linalg.norm(in_quaternion, dim=-1, keepdim=True)
    # normalize quaternion
    in_q = k2 * in_quaternion
    return in_q
//  [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)

                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)

            # 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)

                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)