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

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

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

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

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

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

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

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

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def rotation_matrix(a: TensorType[3], b: TensorType[3]) -> TensorType[3, 3]:
    """Compute the rotation matrix that rotates vector a to vector b.

    Args:
        a: The vector to rotate.
        b: The vector to rotate to.
    Returns:
        The rotation matrix.
    """
    a = a / t.norm(a)
    b = b / t.norm(b)
    v = torch.cross(a, b)
    c = torch.dot(a, b)
    # If vectors are exactly opposite, we add a little noise to one of them
    if c < -1 + 1e-8:
        eps = (torch.rand(3) - 0.5) * 0.01
        return rotation_matrix(a + eps, b)
    s = t.norm(v)
    skew_sym_mat = torch.Tensor(
        [
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0],
        ]
    )
    return torch.eye(3) + skew_sym_mat + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s**2 + 1e-8))
//  [python]
def rotation_matrix(a: TensorType[3], b: TensorType[3]) -> TensorType[3, 3]:
    """Compute the rotation matrix that rotates vector a to vector b.

    Args:
        a: The vector to rotate.
        b: The vector to rotate to.
    Returns:
        The rotation matrix.
    """
    a = a / torch.linalg.norm(a)
    b = b / torch.linalg.norm(b)
    v = torch.cross(a, b)
    c = torch.dot(a, b)
    # If vectors are exactly opposite, we add a little noise to one of them
    if c < -1 + 1e-8:
        eps = (torch.rand(3) - 0.5) * 0.01
        return rotation_matrix(a + eps, b)
    s = torch.linalg.norm(v)
    skew_sym_mat = torch.Tensor(
        [
            [0, -v[2], v[1]],
            [v[2], 0, -v[0]],
            [-v[1], v[0], 0],
        ]
    )
    return torch.eye(3) + skew_sym_mat + skew_sym_mat @ skew_sym_mat * ((1 - c) / (s**2 + 1e-8))
//  [python]
def forward(self, x:torch.Tensor, y:torch.Tensor)-> torch.Tensor:
        norm = get_outnorm(x, self.out_norm)

        if self.kind == 'mat':
            loss = torch.linalg.matrix_norm(
                x - y, ord=self.order).mean()
        else:
            # norm = torch.norm(x - y, p=self.order)
            loss = torch.norm(
                x.view(-1, 1) - y.view(-1, 1), p=self.order)

        return loss*norm
//  [python]
def forward(self, x:torch.Tensor, y:torch.Tensor)-> torch.Tensor:
        norm = get_outnorm(x, self.out_norm)

        if self.kind == 'mat':
            loss = torch.linalg.matrix_norm(
                x - y, ord=self.order).mean()
        else:
            # norm = torch.norm(x - y, p=self.order)
            loss = torch.linalg.norm(
                x.view(-1, 1) - y.view(-1, 1), ord=self.order)

        return loss*norm
//  [python]
def text_embeddings(self, prompts: List[str]) -> torch.Tensor:
        tokens, lens = self.encode_prompts(prompts)
        z_t = self.text_encoder(tokens, lens)
        return z_t / (torch.norm(z_t, dim=-1, keepdim=True) + 1e-12)
//  [python]
def text_embeddings(self, prompts: List[str]) -> torch.Tensor:
        tokens, lens = self.encode_prompts(prompts)
        z_t = self.text_encoder(tokens, lens)
        return z_t / (torch.linalg.norm(z_t, dim=-1, keepdim=True) + 1e-12)
//  [python]
def fc_loss_rot_repr(self, gt_xyz, pred_xyz, mask):
        def to_np_cpu(x):
            return x.detach().cpu().numpy()
        """
        pose_xyz: SMPL batch tensor of shape: [BatchSize, 24, 3, Frames]
        """
        # 'L_Ankle',  # 7, 'R_Ankle',  # 8 , 'L_Foot',  # 10, 'R_Foot',  # 11

        l_ankle_idx, r_ankle_idx = 7, 8
        l_foot_idx, r_foot_idx = 10, 11

        gt_joint_xyz = gt_xyz[:, [l_ankle_idx, l_foot_idx, r_ankle_idx, r_foot_idx], :, :]  # [BatchSize, 4, 3, Frames]
        gt_joint_vel = torch.norm(gt_joint_xyz[:, :, :, 1:] - gt_joint_xyz[:, :, :, :-1], dim=2)  # [BatchSize, 4, Frames]
        fc_mask = (gt_joint_vel <= 0.01)
        pred_joint_xyz = pred_xyz[:, [l_ankle_idx, l_foot_idx, r_ankle_idx, r_foot_idx], :, :]  # [BatchSize, 4, 3, Frames]
        pred_joint_vel = torch.norm(pred_joint_xyz[:, :, :, 1:] - pred_joint_xyz[:, :, :, :-1], dim=2)  # [BatchSize, 4, Frames]
        pred_joint_vel[~fc_mask] = 0  # Blank non-contact velocities frames. [BS,4,FRAMES]
        pred_joint_vel = torch.unsqueeze(pred_joint_vel, dim=2)

        """DEBUG CODE"""
        # print(f'mask: {mask.shape}')
        # print(f'pred_joint_vel: {pred_joint_vel.shape}')
        # plt.title(f'Joint: {joint_idx}')
        # plt.plot(to_np_cpu(gt_joint_vel[0]), label='velocity')
        # plt.plot(to_np_cpu(fc_mask[0]), label='fc')
        # plt.grid()
        # plt.legend()
        # plt.show()
        return self.masked_l2(pred_joint_vel, torch.zeros(pred_joint_vel.shape, device=pred_joint_vel.device),
                              mask[:, :, :, 1:])
//  [python]
def fc_loss_rot_repr(self, gt_xyz, pred_xyz, mask):
        def to_np_cpu(x):
            return x.detach().cpu().numpy()
        """
        pose_xyz: SMPL batch tensor of shape: [BatchSize, 24, 3, Frames]
        """
        # 'L_Ankle',  # 7, 'R_Ankle',  # 8 , 'L_Foot',  # 10, 'R_Foot',  # 11

        l_ankle_idx, r_ankle_idx = 7, 8
        l_foot_idx, r_foot_idx = 10, 11

        gt_joint_xyz = gt_xyz[:, [l_ankle_idx, l_foot_idx, r_ankle_idx, r_foot_idx], :, :]  # [BatchSize, 4, 3, Frames]
        gt_joint_vel = torch.linalg.norm(gt_joint_xyz[:, :, :, 1:] - gt_joint_xyz[:, :, :, :-1], axis=2)  # [BatchSize, 4, Frames]
        fc_mask = (gt_joint_vel <= 0.01)
        pred_joint_xyz = pred_xyz[:, [l_ankle_idx, l_foot_idx, r_ankle_idx, r_foot_idx], :, :]  # [BatchSize, 4, 3, Frames]
        pred_joint_vel = torch.linalg.norm(pred_joint_xyz[:, :, :, 1:] - pred_joint_xyz[:, :, :, :-1], axis=2)  # [BatchSize, 4, Frames]
        pred_joint_vel[~fc_mask] = 0  # Blank non-contact velocities frames. [BS,4,FRAMES]
        pred_joint_vel = torch.unsqueeze(pred_joint_vel, dim=2)

        """DEBUG CODE"""
        # print(f'mask: {mask.shape}')
        # print(f'pred_joint_vel: {pred_joint_vel.shape}')
        # plt.title(f'Joint: {joint_idx}')
        # plt.plot(to_np_cpu(gt_joint_vel[0]), label='velocity')
        # plt.plot(to_np_cpu(fc_mask[0]), label='fc')
        # plt.grid()
        # plt.legend()
        # plt.show()
        return self.masked_l2(pred_joint_vel, torch.zeros(pred_joint_vel.shape, device=pred_joint_vel.device),
                              mask[:, :, :, 1:])
//  [python]
def image_embeddings(self, images: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        z_i = self.image_encoder((images + 1) * 127.5, t)
        return z_i / (torch.norm(z_i, dim=-1, keepdim=True) + 1e-12)
//  [python]
def image_embeddings(self, images: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        z_i = self.image_encoder((images + 1) * 127.5, t)
        return z_i / (torch.linalg.norm(z_i, dim=-1, keepdim=True) + 1e-12)
//  [python]
def check_structure_test(pretrain_file, args1, args2):
    """
    Test that the "copy" method copies the parameters from one model to another

    Also check that the copied models produce the same results
    """
    set_random_seed(1000)
    other = build_model(pretrain_file, *args1)
    other.eval()

    set_random_seed(1001)
    model = build_model(pretrain_file, *args2)
    model.eval()

    assert not torch.allclose(model.delta_embedding.weight, other.delta_embedding.weight)
    assert not torch.allclose(model.output_layers[0].weight, other.output_layers[0].weight)

    model.copy_with_new_structure(other)

    assert torch.allclose(model.delta_embedding.weight, other.delta_embedding.weight)
    assert torch.allclose(model.output_layers[0].weight, other.output_layers[0].weight)
    # the norms will be the same, as the non-zero values are all the same
    assert torch.allclose(torch.norm(model.word_lstm.weight_ih_l0), torch.norm(other.word_lstm.weight_ih_l0))

    # now, check that applying one transition to an initial state
    # results in the same values in the output states for both models
    # as the pattn layer inputs are 0, the output values should be equal
    shift = [parse_transitions.Shift()]
    model_states = test_parse_transitions.build_initial_state(model, 1)
    model_states = model.bulk_apply(model_states, shift)

    other_states = test_parse_transitions.build_initial_state(other, 1)
    other_states = other.bulk_apply(other_states, shift)

    for i, j in zip(other_states[0].word_queue, model_states[0].word_queue):
        assert torch.allclose(i.hx, j.hx, atol=1e-07)
    for i, j in zip(other_states[0].transitions, model_states[0].transitions):
        assert torch.allclose(i.lstm_hx, j.lstm_hx)
        assert torch.allclose(i.lstm_cx, j.lstm_cx)
    for i, j in zip(other_states[0].constituents, model_states[0].constituents):
        assert (i.value is None) == (j.value is None)
        if i.value is not None:
            assert torch.allclose(i.value.tree_hx, j.value.tree_hx, atol=1e-07)
        assert torch.allclose(i.lstm_hx, j.lstm_hx)
        assert torch.allclose(i.lstm_cx, j.lstm_cx)
//  [python]
def check_structure_test(pretrain_file, args1, args2):
    """
    Test that the "copy" method copies the parameters from one model to another

    Also check that the copied models produce the same results
    """
    set_random_seed(1000)
    other = build_model(pretrain_file, *args1)
    other.eval()

    set_random_seed(1001)
    model = build_model(pretrain_file, *args2)
    model.eval()

    assert not torch.allclose(model.delta_embedding.weight, other.delta_embedding.weight)
    assert not torch.allclose(model.output_layers[0].weight, other.output_layers[0].weight)

    model.copy_with_new_structure(other)

    assert torch.allclose(model.delta_embedding.weight, other.delta_embedding.weight)
    assert torch.allclose(model.output_layers[0].weight, other.output_layers[0].weight)
    # the norms will be the same, as the non-zero values are all the same
    assert torch.allclose(torch.linalg.norm(model.word_lstm.weight_ih_l0), torch.linalg.norm(other.word_lstm.weight_ih_l0))

    # now, check that applying one transition to an initial state
    # results in the same values in the output states for both models
    # as the pattn layer inputs are 0, the output values should be equal
    shift = [parse_transitions.Shift()]
    model_states = test_parse_transitions.build_initial_state(model, 1)
    model_states = model.bulk_apply(model_states, shift)

    other_states = test_parse_transitions.build_initial_state(other, 1)
    other_states = other.bulk_apply(other_states, shift)

    for i, j in zip(other_states[0].word_queue, model_states[0].word_queue):
        assert torch.allclose(i.hx, j.hx, atol=1e-07)
    for i, j in zip(other_states[0].transitions, model_states[0].transitions):
        assert torch.allclose(i.lstm_hx, j.lstm_hx)
        assert torch.allclose(i.lstm_cx, j.lstm_cx)
    for i, j in zip(other_states[0].constituents, model_states[0].constituents):
        assert (i.value is None) == (j.value is None)
        if i.value is not None:
            assert torch.allclose(i.value.tree_hx, j.value.tree_hx, atol=1e-07)
        assert torch.allclose(i.lstm_hx, j.lstm_hx)
        assert torch.allclose(i.lstm_cx, j.lstm_cx)
//  [python]
def influential_svd_neurons(
        agg: torch.Tensor, threshold: float = 0.1, norm=1, device=torch.device("cpu")
    ) -> Tuple[sp.coo_matrix, sp.coo_matrix]:
        """
        Returns a dictionary of relative contributions keyed by influential
        SVD neurons for layer up to some threshold

        Parameters
        ----------
        agg : torch.Tensor
            The SVD projections tensor, with some aggregation
            applied. Expected to be 1-D.

        ord (int, float, +/-inf, 'fro', 'nuc', optional)
            order of norm. See
            https://pytorch.org/docs/stable/linalg.html#torch.linalg.norm

        Returns
        -------
        influential_neurons : sp.coo_matrix
            Matrix representing the influential neurons within the threshold
        influential_weights : sp.coo_matrix
            Matrix assigning weights to each influential neuron according to its
            contribution to the threshold
        """
        with torch.no_grad():
            m = torch.norm(
                agg.view((agg.shape[0], -1)), p=norm, dim=1
            ).unsqueeze(0)

            # sort
            ordsmat_vals, ordsmat_indices = torch.sort(m, descending=True)

            # take the cumsum and normalize by total contribution per dim
            cumsum = torch.cumsum(ordsmat_vals, dim=1)
            totalsum = cumsum[:, -1].detach()

            # find the indices within the threshold goal, per dim
            bool_accept = (cumsum / totalsum.unsqueeze(-1)) <= threshold
            accept = torch.sum(bool_accept, dim=1)

            # normalize by final accepted cumsum
            ordsmat_vals /= cumsum[:, accept - 1]

            # add additional accept, ie accept + 1
            try:
                # use range to enumerate over batch size entries of accept
                bool_accept[range(len(accept)), accept] = True
            except IndexError:
                print("taking all values as influential")

            # find accepted synapses, all other values zero.
            # note: it is ordered by largest norm value
            unordered_weights = torch.where(
                bool_accept,
                ordsmat_vals,
                torch.zeros(ordsmat_vals.shape, device=device),
            )
            # re-order to mantain proper neuron ordering
            influential_weights = unordered_weights.gather(
                1, ordsmat_indices.argsort(1)
            )

            influential_neurons = influential_weights.bool().int()

            return matrix_convert(influential_neurons), matrix_convert(
                influential_weights
            )
//  [python]
def influential_svd_neurons(
        agg: torch.Tensor, threshold: float = 0.1, norm=1, device=torch.device("cpu")
    ) -> Tuple[sp.coo_matrix, sp.coo_matrix]:
        """
        Returns a dictionary of relative contributions keyed by influential
        SVD neurons for layer up to some threshold

        Parameters
        ----------
        agg : torch.Tensor
            The SVD projections tensor, with some aggregation
            applied. Expected to be 1-D.

        ord (int, float, +/-inf, 'fro', 'nuc', optional)
            order of norm. See
            https://pytorch.org/docs/stable/linalg.html#torch.linalg.norm

        Returns
        -------
        influential_neurons : sp.coo_matrix
            Matrix representing the influential neurons within the threshold
        influential_weights : sp.coo_matrix
            Matrix assigning weights to each influential neuron according to its
            contribution to the threshold
        """
        with torch.no_grad():
            m = torch.linalg.norm(
                agg.view((agg.shape[0], -1)), ord=norm, dim=1
            ).unsqueeze(0)

            # sort
            ordsmat_vals, ordsmat_indices = torch.sort(m, descending=True)

            # take the cumsum and normalize by total contribution per dim
            cumsum = torch.cumsum(ordsmat_vals, dim=1)
            totalsum = cumsum[:, -1].detach()

            # find the indices within the threshold goal, per dim
            bool_accept = (cumsum / totalsum.unsqueeze(-1)) <= threshold
            accept = torch.sum(bool_accept, dim=1)

            # normalize by final accepted cumsum
            ordsmat_vals /= cumsum[:, accept - 1]

            # add additional accept, ie accept + 1
            try:
                # use range to enumerate over batch size entries of accept
                bool_accept[range(len(accept)), accept] = True
            except IndexError:
                print("taking all values as influential")

            # find accepted synapses, all other values zero.
            # note: it is ordered by largest norm value
            unordered_weights = torch.where(
                bool_accept,
                ordsmat_vals,
                torch.zeros(ordsmat_vals.shape, device=device),
            )
            # re-order to mantain proper neuron ordering
            influential_weights = unordered_weights.gather(
                1, ordsmat_indices.argsort(1)
            )

            influential_neurons = influential_weights.bool().int()

            return matrix_convert(influential_neurons), matrix_convert(
                influential_weights
            )
//  [python]
def get_rotvec(self, eps=1e-6) -> torch.Tensor:
        """
            Return the underlying axis-angle rotation vector.

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

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

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

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

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

        angle = 2 * torch.atan2(
            torch.linalg.norm(quat[..., 1:], dim=-1),
            quat[..., 0]
        )
        
        angle2 = angle * angle
        small_angle_scales = 2 + angle2 / 12 + 7 * angle2 * angle2 / 2880
        large_angle_scales = angle / torch.sin(angle / 2 + eps)
        
        small_angles = (angle <= 1e-3).float()
        rot_vec_scale = small_angle_scales * small_angles + (1 - small_angles) * large_angle_scales
        rot_vec = rot_vec_scale[..., None] * quat[..., 1:]
        return rot_vec
//  [python]
def forward(
        self,
        x,
        global_pe: Optional[torch.Tensor] = None,
    ):
        B, Din, H, W = x.shape
        # Prepare query
        q = torch.broadcast_to(self.query, [x.shape[0], self.hidden_features]).view(B, self.heads,
                                                                                    self.hidden_features // self.heads)
        q = q / (np.sqrt(q.shape[-1]) + 1e-6)
        if self.normalize_q:
            q = q / (torch.norm(q, dim=-1, keepdim=True) + 1e-6)

        qk_score = self._compute_QK_scores(q, x)

        v = self.to_v(x)

        out_per_head_concat = self._compute_attention(qk_score, v, self.rpe)
        out = self.to_out(out_per_head_concat)
        return out
//  [python]
def forward(
        self,
        x,
        global_pe: Optional[torch.Tensor] = None,
    ):
        B, Din, H, W = x.shape
        # Prepare query
        q = torch.broadcast_to(self.query, [x.shape[0], self.hidden_features]).view(B, self.heads,
                                                                                    self.hidden_features // self.heads)
        q = q / (np.sqrt(q.shape[-1]) + 1e-6)
        if self.normalize_q:
            q = q / (torch.linalg.norm(q, dim=-1, keepdim=True) + 1e-6)

        qk_score = self._compute_QK_scores(q, x)

        v = self.to_v(x)

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

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

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

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

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

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

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

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

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def forward(
        self,
        image_embeds: torch.FloatTensor,
        query_embeds: Optional[torch.FloatTensor],
        query_mask: Optional[torch.Tensor],
    ) -> Tuple[torch.FloatTensor]:
        image_class_embeds = self.dense0(image_embeds)
        if query_embeds is None:
            device = image_class_embeds.device
            batch_size, num_patches = image_class_embeds.shape[:2]
            pred_logits = torch.zeros((batch_size, num_patches, self.query_dim)).to(device)
            return (pred_logits, image_class_embeds)

        # Normalize image and text features
        image_class_embeds /= torch.norm(image_class_embeds, dim=-1, keepdim=True) + 1e-6
        query_embeds /= torch.norm(query_embeds, dim=-1, keepdim=True) + 1e-6

        # Get class predictions
        pred_logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds)

        # Apply a learnable shift and scale to logits
        logit_shift = self.logit_shift(image_embeds)
        logit_scale = self.logit_scale(image_embeds)
        logit_scale = self.elu(logit_scale) + 1
        pred_logits = (pred_logits + logit_shift) * logit_scale

        if query_mask is not None:
            if query_mask.ndim > 1:
                query_mask = torch.unsqueeze(query_mask, dim=-2)

            pred_logits = pred_logits.to(torch.float64)
            pred_logits = torch.where(query_mask == 0, -1e6, pred_logits)
            pred_logits = pred_logits.to(torch.float32)

        return (pred_logits, image_class_embeds)
//  [python]
def forward(
        self,
        image_embeds: torch.FloatTensor,
        query_embeds: Optional[torch.FloatTensor],
        query_mask: Optional[torch.Tensor],
    ) -> Tuple[torch.FloatTensor]:
        image_class_embeds = self.dense0(image_embeds)
        if query_embeds is None:
            device = image_class_embeds.device
            batch_size, num_patches = image_class_embeds.shape[:2]
            pred_logits = torch.zeros((batch_size, num_patches, self.query_dim)).to(device)
            return (pred_logits, image_class_embeds)

        # Normalize image and text features
        image_class_embeds /= torch.linalg.norm(image_class_embeds, dim=-1, keepdim=True) + 1e-6
        query_embeds /= torch.linalg.norm(query_embeds, dim=-1, keepdim=True) + 1e-6

        # Get class predictions
        pred_logits = torch.einsum("...pd,...qd->...pq", image_class_embeds, query_embeds)

        # Apply a learnable shift and scale to logits
        logit_shift = self.logit_shift(image_embeds)
        logit_scale = self.logit_scale(image_embeds)
        logit_scale = self.elu(logit_scale) + 1
        pred_logits = (pred_logits + logit_shift) * logit_scale

        if query_mask is not None:
            if query_mask.ndim > 1:
                query_mask = torch.unsqueeze(query_mask, dim=-2)

            pred_logits = pred_logits.to(torch.float64)
            pred_logits = torch.where(query_mask == 0, -1e6, pred_logits)
            pred_logits = pred_logits.to(torch.float32)

        return (pred_logits, image_class_embeds)
//  [python]
def forward(
        self,
        x,
        timesteps: Optional[torch.Tensor] = None,
        global_pe: Optional[torch.Tensor] = None,
    ):
        assert x.shape[2] == self.axial_dim
        B, Din, H, W = x.shape
        if self.timesteps_embed is not None and timesteps is not None:
            q = self.timesteps_embed(timesteps)
            q = q.view(B, self.heads, self.hidden_features // self.heads)
        else:
            # Prepare query
            q = torch.broadcast_to(self.query, [x.shape[0], self.hidden_features]).view(B, self.heads,
                                                                                        self.hidden_features // self.heads)
        if self.normalize_q:
            q = q / (torch.norm(q, dim=-1, keepdim=True) + 1e-6)

        q = q / np.sqrt(q.shape[-1])

        x_with_pos = x if global_pe is None else x + global_pe
        qk_score = self._compute_QK_scores(q, x_with_pos)

        v = self.to_v(x)

        out_per_head_concat = self._compute_attention(qk_score, v, self.rpe)
        out = self.to_out(out_per_head_concat)
        return out
//  [python]
def forward(
        self,
        x,
        timesteps: Optional[torch.Tensor] = None,
        global_pe: Optional[torch.Tensor] = None,
    ):
        assert x.shape[2] == self.axial_dim
        B, Din, H, W = x.shape
        if self.timesteps_embed is not None and timesteps is not None:
            q = self.timesteps_embed(timesteps)
            q = q.view(B, self.heads, self.hidden_features // self.heads)
        else:
            # Prepare query
            q = torch.broadcast_to(self.query, [x.shape[0], self.hidden_features]).view(B, self.heads,
                                                                                        self.hidden_features // self.heads)
        if self.normalize_q:
            q = q / (torch.linalg.norm(q, dim=-1, keepdim=True) + 1e-6)

        q = q / np.sqrt(q.shape[-1])

        x_with_pos = x if global_pe is None else x + global_pe
        qk_score = self._compute_QK_scores(q, x_with_pos)

        v = self.to_v(x)

        out_per_head_concat = self._compute_attention(qk_score, v, self.rpe)
        out = self.to_out(out_per_head_concat)
        return out
//  [python]
def loss(self, target, output):
        N, L = output.shape

        target_reshape = target.squeeze().reshape(L, N)
        fit_term = 0.5 * torch.norm(target_reshape.t() - output, "fro") ** 2

        O = target_reshape.mean(1, keepdims=True)
        reg_term = torch.norm(self.decoder.weight - O, "fro") ** 2

        return fit_term + self.lambd * reg_term
//  [python]
def loss(self, target, output):
        N, L = output.shape

        target_reshape = target.squeeze().reshape(L, N)
        fit_term = 0.5 * torch.linalg.norm(target_reshape.t() - output, "fro") ** 2

        O = target_reshape.mean(1, keepdims=True)
        reg_term = torch.linalg.norm(self.decoder.weight - O, "fro") ** 2

        return fit_term + self.lambd * reg_term
//  [python]
def forward(
        self,
        x,
        timesteps: Optional[torch.Tensor] = None,
        global_pe: Optional[torch.Tensor] = None,
    ):
        assert x.shape[2] == self.axial_dim
        B, Din, H, W = x.shape
        if self.timesteps_embed is not None and timesteps is not None:
            q = self.timesteps_embed(timesteps)
            q = q.view(
                B, self.scale_factor, self.heads, self.hidden_features // self.heads
            )
        else:
            # Prepare query
            # assert False, "Don't unset timesteps unless explicitly done for testing"
            q = torch.broadcast_to(
                self.query, [x.shape[0], self.scale_factor * self.hidden_features]
            ).view(
                B, self.scale_factor, self.heads, self.hidden_features // self.heads
            )
        if self.normalize_q:
            q = q / (torch.norm(q, dim=-1, keepdim=True) + 1e-6)
        q = q / np.sqrt(q.shape[-1])

        qk_score = self._compute_QK_scores(q, x)

        v = self.to_v(x)

        out_per_head_concat = self._compute_attention(qk_score, v, self.rpe)
        out = self.to_out(out_per_head_concat)
        return out
//  [python]
def forward(
        self,
        x,
        timesteps: Optional[torch.Tensor] = None,
        global_pe: Optional[torch.Tensor] = None,
    ):
        assert x.shape[2] == self.axial_dim
        B, Din, H, W = x.shape
        if self.timesteps_embed is not None and timesteps is not None:
            q = self.timesteps_embed(timesteps)
            q = q.view(
                B, self.scale_factor, self.heads, self.hidden_features // self.heads
            )
        else:
            # Prepare query
            # assert False, "Don't unset timesteps unless explicitly done for testing"
            q = torch.broadcast_to(
                self.query, [x.shape[0], self.scale_factor * self.hidden_features]
            ).view(
                B, self.scale_factor, self.heads, self.hidden_features // self.heads
            )
        if self.normalize_q:
            q = q / (torch.linalg.norm(q, dim=-1, keepdim=True) + 1e-6)
        q = q / np.sqrt(q.shape[-1])

        qk_score = self._compute_QK_scores(q, x)

        v = self.to_v(x)

        out_per_head_concat = self._compute_attention(qk_score, v, self.rpe)
        out = self.to_out(out_per_head_concat)
        return out
//  [python]
def _get_volume(self, params=Points.empty(), device='cpu'):
        _, _, _, dir_1, dir_2 = self.domain._construct_parallelogram(params, device=device)
        side_length1 = torch.norm(dir_1, dim=1)
        side_length2 = torch.norm(dir_2, dim=1)
        return 2 * (side_length1 + side_length2).reshape(-1, 1)
//  [python]
def _get_volume(self, params=Points.empty(), device='cpu'):
        _, _, _, dir_1, dir_2 = self.domain._construct_parallelogram(params, device=device)
        side_length1 = torch.linalg.norm(dir_1, dim=1)
        side_length2 = torch.linalg.norm(dir_2, dim=1)
        return 2 * (side_length1 + side_length2).reshape(-1, 1)
//  [python]
def pt_render_poses(datadir: str, n_frames: int, frame_h: int = 800, frame_w: int = 800):
    scene = PhototourismScenes.get_scene_from_datadir(datadir)
    idx, _ = get_ids_for_split(datadir, split='train')
    train_poses, kinvs, bounds, res = load_camera_metadata(datadir, idx)
    train_poses, kinvs, bounds = scale_cam_metadata(train_poses, kinvs, bounds, scale=0.05)

    # build camera intrinsic from appropriate focal distance and cx, cy. Good for trevi
    k = np.array([[780.0, 0, -frame_w / 2], [0, -780, -frame_h / 2], [0, 0, -1]])
    kinv = torch.from_numpy(np.linalg.inv(k)).to(torch.float32)

    bounds = torch.from_numpy(bounds).float()
    train_poses = torch.from_numpy(train_poses).float()

    r_poses = pt_spiral_path(scene, train_poses, n_frames=n_frames, zrate=1, n_rots=1)

    all_rays_o, all_rays_d, camera_ids, near_fars = [], [], [], []
    for pose_id, pose in enumerate(r_poses):
        pose = pose.float()
        rays_o, rays_d = get_rays_tourism(frame_h, frame_w, kinv, pose)
        rays_o = rays_o.view(-1, 3)
        rays_d = rays_d.view(-1, 3)
        all_rays_o.append(rays_o)
        all_rays_d.append(rays_d)

        # Find the closest cam
        closest_cam_idx = torch.norm(
            train_poses[:, :3, :].view(train_poses.shape[0], -1) - pose.view(-1), dim=1
        ).argmin()

        # For brandenburg and trevi
        if scene == PhototourismScenes.BRANDENBURG or scene == PhototourismScenes.TREVI:
            near_fars.append((
                bounds[closest_cam_idx] + torch.tensor([0.01, 0.0])
            ))
        elif scene == PhototourismScenes.SACRE:
            near_fars.append((
                bounds[closest_cam_idx] + torch.tensor([0.05, 0.0])
            ))
    # camera-IDs. They are floats interpolating between 2 appearance embeddings.
    x = torch.linspace(-1, 1, len(r_poses))
    s = 0.3
    camera_ids = 1/(s * math.sqrt(2 * torch.pi)) * torch.exp(- (x)**2 / (2 * s**2))
    camera_ids = (camera_ids - camera_ids.min()) / camera_ids.max()
    all_rays_o = torch.stack(all_rays_o, dim=0)
    all_rays_d = torch.stack(all_rays_d, dim=0)
    near_fars = torch.stack(near_fars, dim=0)
    return all_rays_o, all_rays_d, camera_ids, near_fars
//  [python]
def pt_render_poses(datadir: str, n_frames: int, frame_h: int = 800, frame_w: int = 800):
    scene = PhototourismScenes.get_scene_from_datadir(datadir)
    idx, _ = get_ids_for_split(datadir, split='train')
    train_poses, kinvs, bounds, res = load_camera_metadata(datadir, idx)
    train_poses, kinvs, bounds = scale_cam_metadata(train_poses, kinvs, bounds, scale=0.05)

    # build camera intrinsic from appropriate focal distance and cx, cy. Good for trevi
    k = np.array([[780.0, 0, -frame_w / 2], [0, -780, -frame_h / 2], [0, 0, -1]])
    kinv = torch.from_numpy(np.linalg.inv(k)).to(torch.float32)

    bounds = torch.from_numpy(bounds).float()
    train_poses = torch.from_numpy(train_poses).float()

    r_poses = pt_spiral_path(scene, train_poses, n_frames=n_frames, zrate=1, n_rots=1)

    all_rays_o, all_rays_d, camera_ids, near_fars = [], [], [], []
    for pose_id, pose in enumerate(r_poses):
        pose = pose.float()
        rays_o, rays_d = get_rays_tourism(frame_h, frame_w, kinv, pose)
        rays_o = rays_o.view(-1, 3)
        rays_d = rays_d.view(-1, 3)
        all_rays_o.append(rays_o)
        all_rays_d.append(rays_d)

        # Find the closest cam
        closest_cam_idx = torch.linalg.norm(
            train_poses[:, :3, :].view(train_poses.shape[0], -1) - pose.view(-1), dim=1
        ).argmin()

        # For brandenburg and trevi
        if scene == PhototourismScenes.BRANDENBURG or scene == PhototourismScenes.TREVI:
            near_fars.append((
                bounds[closest_cam_idx] + torch.tensor([0.01, 0.0])
            ))
        elif scene == PhototourismScenes.SACRE:
            near_fars.append((
                bounds[closest_cam_idx] + torch.tensor([0.05, 0.0])
            ))
    # camera-IDs. They are floats interpolating between 2 appearance embeddings.
    x = torch.linspace(-1, 1, len(r_poses))
    s = 0.3
    camera_ids = 1/(s * math.sqrt(2 * torch.pi)) * torch.exp(- (x)**2 / (2 * s**2))
    camera_ids = (camera_ids - camera_ids.min()) / camera_ids.max()
    all_rays_o = torch.stack(all_rays_o, dim=0)
    all_rays_d = torch.stack(all_rays_d, dim=0)
    near_fars = torch.stack(near_fars, dim=0)
    return all_rays_o, all_rays_d, camera_ids, near_fars
//  [python]
def _compute_side_length(self, dir_1, dir_2):
        # essentially computes the same as volume, but we need to set the view 
        # that we can use the computes values for different cases
        side_length1 = torch.norm(dir_1, dim=1).view(-1, 1, 1)
        side_length2 = torch.norm(dir_2, dim=1).view(-1, 1, 1)
        total_length = 2 * (side_length1 + side_length2)
        return side_length1,side_length2,total_length
//  [python]
def _compute_side_length(self, dir_1, dir_2):
        # essentially computes the same as volume, but we need to set the view 
        # that we can use the computes values for different cases
        side_length1 = torch.linalg.norm(dir_1, dim=1).view(-1, 1, 1)
        side_length2 = torch.linalg.norm(dir_2, dim=1).view(-1, 1, 1)
        total_length = 2 * (side_length1 + side_length2)
        return side_length1,side_length2,total_length
//  [python]
def _get_normal_direction(self, direction, device):
        # to get normal vector in 2d switch x and y coordinate and multiply 
        # one coordinate with -1
        normal = torch.index_select(direction, 1,
                                    torch.tensor([1, 0], device=device))
        normal[:, :1] *= -1
        return torch.divide(normal, torch.norm(normal, dim=1).reshape(-1, 1))
//  [python]
def _get_normal_direction(self, direction, device):
        # to get normal vector in 2d switch x and y coordinate and multiply 
        # one coordinate with -1
        normal = torch.index_select(direction, 1,
                                    torch.tensor([1, 0], device=device))
        normal[:, :1] *= -1
        return torch.divide(normal, torch.linalg.norm(normal, dim=1).reshape(-1, 1))
//  [python]
def _compute_barycentric_grid(self, n, dir_1, dir_2, device):
        side_length_1 = torch.norm(dir_1, dim=1)
        side_length_2 = torch.norm(dir_2, dim=1)
        # scale the number of point w.r.t. the 'form' of the parallelogram
        n_1 = int(torch.sqrt(n*side_length_1/side_length_2))
        n_2 = int(torch.sqrt(n*side_length_2/side_length_1))
        x = torch.linspace(0, 1, n_1+2, device=device)[1:-1] # create inner grid, so remove
        y = torch.linspace(0, 1, n_2+2, device=device)[1:-1] # first and last value
        bary_coords = torch.stack(torch.meshgrid((x, y))).T.reshape(-1, 2)
        return bary_coords
//  [python]
def _compute_barycentric_grid(self, n, dir_1, dir_2, device):
        side_length_1 = torch.linalg.norm(dir_1, dim=1)
        side_length_2 = torch.linalg.norm(dir_2, dim=1)
        # scale the number of point w.r.t. the 'form' of the parallelogram
        n_1 = int(torch.sqrt(n*side_length_1/side_length_2))
        n_2 = int(torch.sqrt(n*side_length_2/side_length_1))
        x = torch.linspace(0, 1, n_1+2, device=device)[1:-1] # create inner grid, so remove
        y = torch.linspace(0, 1, n_2+2, device=device)[1:-1] # first and last value
        bary_coords = torch.stack(torch.meshgrid((x, y))).T.reshape(-1, 2)
        return bary_coords
//  [python]
def __init__(
        self,
        rot_mats: Optional[torch.Tensor] = None,
        quats: Optional[torch.Tensor] = None,
        normalize_quats: bool = True,
    ):
        """
        Args:
            rot_mats:
                A [*, 3, 3] rotation matrix tensor. Mutually exclusive with
                quats
            quats:
                A [*, 4] quaternion. Mutually exclusive with rot_mats. If
                normalize_quats is not True, must be a unit quaternion
            normalize_quats:
                If quats is specified, whether to normalize quats
        """
        if (rot_mats is None and quats is None) or (
            rot_mats is not None and quats is not None
        ):
            raise ValueError("Exactly one input argument must be specified")

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

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

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

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

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

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

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

        self._rot_mats = rot_mats
        self._quats = quats
//  [python]
def normal(self, points, params=Points.empty(), device='cpu'):
        points, params, device = \
            self._transform_input_for_normals(points, params, device)
        origin, _, _, dir_1, dir_2 = \
            self.domain._construct_parallelogram(points.join(params), device)
        points = points[:, list(self.space.keys())].as_tensor
        normals = torch.zeros_like(points, device=device)
        bary_x, bary_y = self.domain._solve_lgs(points - origin, dir_1, dir_2)
        normal_dir_1 = self._get_normal_direction(dir_1, device)
        normal_dir_2 = -self._get_normal_direction(dir_2, device)
        # compute for each point what the normal vector should be, by checking the
        # value of the local barycentric coordinate = 0 or 1
        self._add_local_normal_vector(normals, bary_x, bary_y, normal_dir_1,
                                      normal_dir_2, 0.0)
        self._add_local_normal_vector(normals, bary_x, bary_y, normal_dir_1,
                                      normal_dir_2, 1.0)
        # scale normal vectors if there where in a corner:
        return torch.divide(normals, torch.norm(normals, dim=1).reshape(-1, 1))
//  [python]
def normal(self, points, params=Points.empty(), device='cpu'):
        points, params, device = \
            self._transform_input_for_normals(points, params, device)
        origin, _, _, dir_1, dir_2 = \
            self.domain._construct_parallelogram(points.join(params), device)
        points = points[:, list(self.space.keys())].as_tensor
        normals = torch.zeros_like(points, device=device)
        bary_x, bary_y = self.domain._solve_lgs(points - origin, dir_1, dir_2)
        normal_dir_1 = self._get_normal_direction(dir_1, device)
        normal_dir_2 = -self._get_normal_direction(dir_2, device)
        # compute for each point what the normal vector should be, by checking the
        # value of the local barycentric coordinate = 0 or 1
        self._add_local_normal_vector(normals, bary_x, bary_y, normal_dir_1,
                                      normal_dir_2, 0.0)
        self._add_local_normal_vector(normals, bary_x, bary_y, normal_dir_1,
                                      normal_dir_2, 1.0)
        # scale normal vectors if there where in a corner:
        return torch.divide(normals, torch.linalg.norm(normals, dim=1).reshape(-1, 1))
//  [python]
def get_norm(self, x):
        norm = torch.norm(x, dim=-1, keepdim=True)
        norm[norm == 0] = 1
        
        assert norm.shape[-1] == 1
        assert norm.shape[0] == x.shape[0]
        assert norm.shape[1] == x.shape[1]
        assert x.shape[1] == 1

        return norm
//  [python]
def get_norm(self, x):
        norm = torch.linalg.norm(x, dim=-1, keepdim=True)
        norm[norm == 0] = 1
        
        assert norm.shape[-1] == 1
        assert norm.shape[0] == x.shape[0]
        assert norm.shape[1] == x.shape[1]
        assert x.shape[1] == 1

        return norm
//  [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.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})