jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def orthogonal_(tensor, gain=1):
    # proper orthogonal init, see https://github.com/pytorch/pytorch/pull/10672
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor[0].numel()
    flattened = tensor.new_empty(rows, cols).normal_(0, 1)

    for i in range(0, rows, cols):
        # Compute the qr factorization
        q, r = torch.linalg.qr(flattened[i:i + cols].t())
        # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
        q *= torch.diag(r, 0).sign()
        q.t_()

        with torch.no_grad():
            tensor[i:i + cols].view_as(q).copy_(q)

    with torch.no_grad():
        tensor.mul_(gain)
    return tensor
//  [python]
def orthogonal_(tensor, gain=1):
    r"""Fills the input `Tensor` with a (semi) orthogonal matrix, as
    described in "Exact solutions to the nonlinear dynamics of learning in deep
    linear neural networks" - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the
    trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional `torch.Tensor`, where :math:`n \geq 2`
        gain: optional scaling factor

    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.orthogonal_(w)
    """
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor[0].numel()
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def orthogonal_(tensor, gain=1):
    r"""Fills the input `Tensor` with a (semi) orthogonal matrix, as
    described in "Exact solutions to the nonlinear dynamics of learning in deep
    linear neural networks" - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the
    trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional `torch.Tensor`, where :math:`n \geq 2`
        gain: optional scaling factor

    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.orthogonal_(w)
    """
    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor[0].numel()
    flattened = tensor.new(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.linalg.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def _get_weight(self):
        weight = self.scale * self.weight  # (cout, cin)
        if self.orthogonalize:
            # qr decomposition also normalize the bases
            q, r = torch.qr(weight)
            d = torch.diag(r, 0)
            ph = d.sign()  # +1 or -1   # loss gradient to r

            if self.normalize:
                weight = q * ph  # (cout, cin)
            else:
                ori_norm = weight.norm(p=2, dim=0, keepdim=True)  # (1, cin)
                weight = q * ph  # (cout, cin)
                weight = weight * ori_norm

        elif self.normalize:
            demod = torch.rsqrt(weight.pow(2).sum(dim=0, keepdim=True) + self.eps)  # (1, cin)
            weight = weight * demod

        return weight
//  [python]
def _get_weight(self):
        weight = self.scale * self.weight  # (cout, cin)
        if self.orthogonalize:
            # qr decomposition also normalize the bases
            q, r = torch.linalg.qr(weight)
            d = torch.diag(r, 0)
            ph = d.sign()  # +1 or -1   # loss gradient to r

            if self.normalize:
                weight = q * ph  # (cout, cin)
            else:
                ori_norm = weight.norm(p=2, dim=0, keepdim=True)  # (1, cin)
                weight = q * ph  # (cout, cin)
                weight = weight * ori_norm

        elif self.normalize:
            demod = torch.rsqrt(weight.pow(2).sum(dim=0, keepdim=True) + self.eps)  # (1, cin)
            weight = weight * demod

        return weight
//  [python]
def test_ormqr(self):
        mat1 = torch.randn(10, 10)
        mat2 = torch.randn(10, 10)
        q, r = torch.qr(mat1)
        m, tau = torch.geqrf(mat1)

        res1 = torch.mm(q, mat2)
        res2, _ = torch.ormqr(m, tau, mat2)
        self.assertEqual(res1, res2)

        res1 = torch.mm(mat2, q)
        res2, _ = torch.ormqr(m, tau, mat2, False)
        self.assertEqual(res1, res2)

        res1 = torch.mm(q.t(), mat2)
        res2, _ = torch.ormqr(m, tau, mat2, True, True)
        self.assertEqual(res1, res2)

        res1 = torch.mm(mat2, q.t())
        res2, _ = torch.ormqr(m, tau, mat2, False, True)
        self.assertEqual(res1, res2)
//  [python]
def test_ormqr(self):
        mat1 = torch.randn(10, 10)
        mat2 = torch.randn(10, 10)
        q, r = torch.linalg.qr(mat1)
        m, tau = torch.geqrf(mat1)

        res1 = torch.mm(q, mat2)
        res2, _ = torch.ormqr(m, tau, mat2)
        self.assertEqual(res1, res2)

        res1 = torch.mm(mat2, q)
        res2, _ = torch.ormqr(m, tau, mat2, False)
        self.assertEqual(res1, res2)

        res1 = torch.mm(q.t(), mat2)
        res2, _ = torch.ormqr(m, tau, mat2, True, True)
        self.assertEqual(res1, res2)

        res1 = torch.mm(mat2, q.t())
        res2, _ = torch.ormqr(m, tau, mat2, False, True)
        self.assertEqual(res1, res2)
//  [python]
def rand(self):
        """
        Generate random positive definite point
        """
        if self._k == 1:
            d = torch.ones(self._n, 1) + torch.rand(self._n, 1)
            X = torch.randn(self._n, self._n)
            q, _ = torch.qr(X)
        else:
            d = torch.ones(self._k, self._n, 1) + torch.rand(self._k, self._n, 1)
            X = torch.zeros(self._k, self._n, self._n)
            for i in range(self._k):
                X[i], _ = torch.qr(torch.randn(self._n, self._n))

        return multiprod(X, d * multitransp(X))
//  [python]
def rand(self):
        """
        Generate random positive definite point
        """
        if self._k == 1:
            d = torch.ones(self._n, 1) + torch.rand(self._n, 1)
            X = torch.randn(self._n, self._n)
            q, _ = torch.linalg.qr(X)
        else:
            d = torch.ones(self._k, self._n, 1) + torch.rand(self._k, self._n, 1)
            X = torch.zeros(self._k, self._n, self._n)
            for i in range(self._k):
                X[i], _ = torch.linalg.qr(torch.randn(self._n, self._n))

        return multiprod(X, d * multitransp(X))
//  [python]
def test_fisher_matrix_matrix_matmul(self):
        model = torch.nn.Sequential(
            torch.nn.Linear(1, 400),
            torch.nn.ELU(),
            torch.nn.Linear(400, 400),
            torch.nn.ELU(),
            torch.nn.Linear(400, 1),
        )

        data = torch.randn(1500, 1)

        fvp = FVPR_FD(model, data)

        numpars = 0
        for p in model.parameters():
            numpars += p.numel()

        orthmat, _ = torch.qr(torch.randn(numpars, 80))
        emat = 1e-2 * torch.randn(80, 2)

        full_matmul = fvp.matmul(orthmat @ emat)
        split_matmul = fvp.matmul(orthmat) @ emat

        # check that F (Vy) = FV y
        self.assertLess(
            torch.norm(full_matmul - split_matmul) / split_matmul.norm(), 1e-2
        )

        # check that matrix columns work
        self.assertLess(
            torch.norm(full_matmul[:, 0] - fvp.matmul(orthmat @ emat[:, 0])), 1e-5
        )
//  [python]
def test_fisher_matrix_matrix_matmul(self):
        model = torch.nn.Sequential(
            torch.nn.Linear(1, 400),
            torch.nn.ELU(),
            torch.nn.Linear(400, 400),
            torch.nn.ELU(),
            torch.nn.Linear(400, 1),
        )

        data = torch.randn(1500, 1)

        fvp = FVPR_FD(model, data)

        numpars = 0
        for p in model.parameters():
            numpars += p.numel()

        orthmat, _ = torch.linalg.qr(torch.randn(numpars, 80))
        emat = 1e-2 * torch.randn(80, 2)

        full_matmul = fvp.matmul(orthmat @ emat)
        split_matmul = fvp.matmul(orthmat) @ emat

        # check that F (Vy) = FV y
        self.assertLess(
            torch.norm(full_matmul - split_matmul) / split_matmul.norm(), 1e-2
        )

        # check that matrix columns work
        self.assertLess(
            torch.norm(full_matmul[:, 0] - fvp.matmul(orthmat @ emat[:, 0])), 1e-5
        )
//  [python]
def fit(self, X: torch.Tensor):
        """
        Generate random coefficients from the shape of the training data
        Args:
            X: temsor of size (n_samples, n_features)
        """
        _, n_data_features = X.shape
        size = (n_data_features, n_data_features)
        # We compute random features in orthogonal blocks of size n_data_features
        n_stacks = int(np.ceil(self.rff_n_features / n_data_features))
        # As a result of orthogonality constraint we can have slightly more random features than specified
        rff_n_features = n_stacks * n_data_features

        random_weights_ = []
        for _ in range(n_stacks):
            # Iterate over stacks, building feature coefficients orthogonal in the dimension of the observations
            W = torch.randn(size, dtype=X.dtype, device=X.device)
            Q, _ = torch.qr(W, some=False)

            chi2 = (
                dist.Chi2(df=torch.tensor([n_data_features], dtype=X.dtype, device=X.device))
                .sample((n_data_features,))
                .sqrt()
                .squeeze()
            )
            if chi2.numel() > 1:
                chi2 = chi2.diag()
            else:
                chi2 = chi2.unsqueeze(0)
            SQ = Q @ chi2  # size (n_data_features, n_data_features)
            random_weights_ += [SQ]

        self.random_weights_ = torch.vstack(random_weights_).T  # Shape (n_data_features, rff_n_features)

        if isinstance(self.lengthscale, list) or isinstance(self.lengthscale, tuple):
            # Marginalise lengthscale over some uniform prior to provide a more flexible function class
            lengthscale = (
                torch.rand(1, n_stacks, dtype=X.dtype, device=X.device) * (self.lengthscale[1] - self.lengthscale[0])
                + self.lengthscale[0]
            )
            lengthscale = lengthscale.repeat_interleave(n_data_features).unsqueeze(0)
        else:
            # Scale feature coefficients by fixed lengthscale
            lengthscale = self.lengthscale

        self.random_weights_ /= lengthscale
        # Generate random bias term
        self.random_offset_ = torch.rand(rff_n_features, dtype=X.dtype, device=X.device) * np.pi * 2

        self.fitted = True
//  [python]
def fit(self, X: torch.Tensor):
        """
        Generate random coefficients from the shape of the training data
        Args:
            X: temsor of size (n_samples, n_features)
        """
        _, n_data_features = X.shape
        size = (n_data_features, n_data_features)
        # We compute random features in orthogonal blocks of size n_data_features
        n_stacks = int(np.ceil(self.rff_n_features / n_data_features))
        # As a result of orthogonality constraint we can have slightly more random features than specified
        rff_n_features = n_stacks * n_data_features

        random_weights_ = []
        for _ in range(n_stacks):
            # Iterate over stacks, building feature coefficients orthogonal in the dimension of the observations
            W = torch.randn(size, dtype=X.dtype, device=X.device)
            Q, _ = torch.linalg.qr(W, mode='reduced')

            chi2 = (
                dist.Chi2(df=torch.tensor([n_data_features], dtype=X.dtype, device=X.device))
                .sample((n_data_features,))
                .sqrt()
                .squeeze()
            )
            if chi2.numel() > 1:
                chi2 = chi2.diag()
            else:
                chi2 = chi2.unsqueeze(0)
            SQ = Q @ chi2  # size (n_data_features, n_data_features)
            random_weights_ += [SQ]

        self.random_weights_ = torch.vstack(random_weights_).T  # Shape (n_data_features, rff_n_features)

        if isinstance(self.lengthscale, list) or isinstance(self.lengthscale, tuple):
            # Marginalise lengthscale over some uniform prior to provide a more flexible function class
            lengthscale = (
                torch.rand(1, n_stacks, dtype=X.dtype, device=X.device) * (self.lengthscale[1] - self.lengthscale[0])
                + self.lengthscale[0]
            )
            lengthscale = lengthscale.repeat_interleave(n_data_features).unsqueeze(0)
        else:
            # Scale feature coefficients by fixed lengthscale
            lengthscale = self.lengthscale

        self.random_weights_ /= lengthscale
        # Generate random bias term
        self.random_offset_ = torch.rand(rff_n_features, dtype=X.dtype, device=X.device) * np.pi * 2

        self.fitted = True
//  [python]
def orthogonal_matrix_chunk(cols, device = None):
    unstructured_block = torch.randn((cols, cols), device = device)
    q, r = torch.qr(unstructured_block.cpu(), some = True)
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()
//  [python]
def orthogonal_matrix_chunk(cols, device = None):
    unstructured_block = torch.randn((cols, cols), device = device)
    q, r = torch.linalg.qr(unstructured_block.cpu())
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()
//  [python]
def unify_token(self, token_feature):
        """
            Unify Token Representation
        """
        window_size = self.context_window_size

        alpha_alignment = torch.zeros(token_feature.size()[0], device=token_feature.device)
        alpha_novelty = torch.zeros(token_feature.size()[0], device=token_feature.device)

        for k in range(token_feature.size()[0]):
            left_window = token_feature[k - window_size:k, :]
            right_window = token_feature[k + 1:k + window_size + 1, :]
            window_matrix = torch.cat([left_window, right_window, token_feature[k, :][None, :]])
            Q, R = torch.qr(window_matrix.T)

            r = R[:, -1]
            alpha_alignment[k] = torch.mean(self.norm_vector(R[:-1, :-1], dim=0), dim=1).matmul(R[:-1, -1]) / torch.norm(r[:-1])
            alpha_alignment[k] = 1 / (alpha_alignment[k] * window_matrix.size()[0] * 2)
            alpha_novelty[k] = torch.abs(r[-1]) / torch.norm(r)

        # Sum Norm
        alpha_alignment = alpha_alignment / torch.sum(alpha_alignment)  # Normalization Choice
        alpha_novelty = alpha_novelty / torch.sum(alpha_novelty)

        alpha = alpha_novelty + alpha_alignment
        alpha = alpha / torch.sum(alpha)  # Normalize

        out_embedding = torch.mv(token_feature.t(), alpha)
        return out_embedding
//  [python]
def unify_token(self, token_feature):
        """
            Unify Token Representation
        """
        window_size = self.context_window_size

        alpha_alignment = torch.zeros(token_feature.size()[0], device=token_feature.device)
        alpha_novelty = torch.zeros(token_feature.size()[0], device=token_feature.device)

        for k in range(token_feature.size()[0]):
            left_window = token_feature[k - window_size:k, :]
            right_window = token_feature[k + 1:k + window_size + 1, :]
            window_matrix = torch.cat([left_window, right_window, token_feature[k, :][None, :]])
            Q, R = torch.linalg.qr(window_matrix.T)

            r = R[:, -1]
            alpha_alignment[k] = torch.mean(self.norm_vector(R[:-1, :-1], dim=0), dim=1).matmul(R[:-1, -1]) / torch.norm(r[:-1])
            alpha_alignment[k] = 1 / (alpha_alignment[k] * window_matrix.size()[0] * 2)
            alpha_novelty[k] = torch.abs(r[-1]) / torch.norm(r)

        # Sum Norm
        alpha_alignment = alpha_alignment / torch.sum(alpha_alignment)  # Normalization Choice
        alpha_novelty = alpha_novelty / torch.sum(alpha_novelty)

        alpha = alpha_novelty + alpha_alignment
        alpha = alpha / torch.sum(alpha)  # Normalize

        out_embedding = torch.mv(token_feature.t(), alpha)
        return out_embedding
//  [python]
def __init__(self, in_channel, out_channel=None):
        super().__init__()
        
        if out_channel is None:
            out_channel = in_channel
        weight = torch.randn(in_channel, out_channel)
        q, _ = torch.qr(weight)
        weight = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(weight)
//  [python]
def __init__(self, in_channel, out_channel=None):
        super().__init__()
        
        if out_channel is None:
            out_channel = in_channel
        weight = torch.randn(in_channel, out_channel)
        q, _ = torch.linalg.qr(weight)
        weight = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(weight)
//  [python]
def orthogonal_matrix_chunk(cols, qr_uniform_q = False, device = None):
    unstructured_block = torch.randn((cols, cols), device = device)
    q, r = torch.qr(unstructured_block.cpu(), some = True)
    q, r = map(lambda t: t.to(device), (q, r))

    # proposed by @Parskatt
    # to make sure Q is uniform https://arxiv.org/pdf/math-ph/0609050.pdf
    if qr_uniform_q:
        d = torch.diag(r, 0)
        q *= d.sign()
    return q.t()
//  [python]
def orthogonal_matrix_chunk(cols, qr_uniform_q = False, device = None):
    unstructured_block = torch.randn((cols, cols), device = device)
    q, r = torch.linalg.qr(unstructured_block.cpu())
    q, r = map(lambda t: t.to(device), (q, r))

    # proposed by @Parskatt
    # to make sure Q is uniform https://arxiv.org/pdf/math-ph/0609050.pdf
    if qr_uniform_q:
        d = torch.diag(r, 0)
        q *= d.sign()
    return q.t()
//  [python]
def __init__(self, n_channels=3, lu_factorize=False):
        super().__init__()
        self.lu_factorize = lu_factorize
        w = torch.randn(n_channels, n_channels)
        w = torch.qr(w)[0]
        if lu_factorize:
            p, l, u = torch.btriunpack(*w.unsqueeze(0).btrifact())
            self.p, self.l, self.u = nn.Parameter(p.squeeze()), nn.Parameter(l.squeeze()), nn.Parameter(u.squeeze())
            s = self.u.diag()
            self.log_s = nn.Parameter(s.abs().log())
            self.register_buffer('sign_s', s.sign())
            self.register_buffer('l_mask', torch.tril(torch.ones_like(self.l), -1))
        else:
            self.w = nn.Parameter(w)
//  [python]
def __init__(self, n_channels=3, lu_factorize=False):
        super().__init__()
        self.lu_factorize = lu_factorize
        w = torch.randn(n_channels, n_channels)
        w = torch.linalg.qr(w)[0]
        if lu_factorize:
            p, l, u = torch.btriunpack(*w.unsqueeze(0).btrifact())
            self.p, self.l, self.u = nn.Parameter(p.squeeze()), nn.Parameter(l.squeeze()), nn.Parameter(u.squeeze())
            s = self.u.diag()
            self.log_s = nn.Parameter(s.abs().log())
            self.register_buffer('sign_s', s.sign())
            self.register_buffer('l_mask', torch.tril(torch.ones_like(self.l), -1))
        else:
            self.w = nn.Parameter(w)
//  [python]
def hutchpp(linear_operator, dimension, n_queries):
    A = linear_operator
    d = dimension
    m = n_queries
    S = torch.randn(d, m // 3, device=A.device)
    G = torch.randn(d, m // 3, device=A.device)
    Q, _ = torch.qr(A.matvec(S))
    proj = G - Q @ (Q.T @ G)
    return torch.trace(Q.T @ A.matvec(Q)) + (3.0 / m) * torch.trace(
        proj.T @ A.matvec(proj)
    )
//  [python]
def hutchpp(linear_operator, dimension, n_queries):
    A = linear_operator
    d = dimension
    m = n_queries
    S = torch.randn(d, m // 3, device=A.device)
    G = torch.randn(d, m // 3, device=A.device)
    Q, _ = torch.linalg.qr(A.matvec(S))
    proj = G - Q @ (Q.T @ G)
    return torch.trace(Q.T @ A.matvec(Q)) + (3.0 / m) * torch.trace(
        proj.T @ A.matvec(proj)
    )
//  [python]
def __init__(self, c):
        super(Invertible1x1Conv, self).__init__()
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0, bias=False)

        # Sample a random orthonormal matrix to initialize weights
        W = torch.qr(torch.FloatTensor(c, c).normal_())[0]

        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:,0] = -1*W[:,0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W
//  [python]
def __init__(self, c):
        super(Invertible1x1Conv, self).__init__()
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0, bias=False)

        # Sample a random orthonormal matrix to initialize weights
        W = torch.linalg.qr(torch.FloatTensor(c, c).normal_())[0]

        # Ensure determinant is 1.0 not -1.0
        if torch.det(W) < 0:
            W[:,0] = -1*W[:,0]
        W = W.view(c, c, 1)
        self.conv.weight.data = W
//  [python]
def __init__(self, dim_inv, dim_spu, n_envs):
        super().__init__(dim_inv, dim_spu, n_envs)

        self.scramble, _ = torch.qr(torch.randn(self.dim, self.dim))
//  [python]
def __init__(self, dim_inv, dim_spu, n_envs):
        super().__init__(dim_inv, dim_spu, n_envs)

        self.scramble, _ = torch.linalg.qr(torch.randn(self.dim, self.dim))
//  [python]
def genOrthgonal(dim):
    a = torch.zeros((dim, dim)).normal_(0, 1)
    q, r = torch.qr(a)
    d = torch.diag(r, 0).sign()
    diag_size = d.size(0)
    d_exp = d.view(1, diag_size).expand(diag_size, diag_size)
    q.mul_(d_exp)
    return q
//  [python]
def genOrthgonal(dim):
    a = torch.zeros((dim, dim)).normal_(0, 1)
    q, r = torch.linalg.qr(a)
    d = torch.diag(r, 0).sign()
    diag_size = d.size(0)
    d_exp = d.view(1, diag_size).expand(diag_size, diag_size)
    q.mul_(d_exp)
    return q
//  [python]
def eigsh_func(f, dtype, device, n, k=6, n_iter=4, L=None):
    r"""Top-k Eigendecomposition of PSD linear operator :math:`f(x) = Ax`.

    Provides a way to use :meth:`fastpca.eigsh` when the matrix
    to eigendecompose is only accessible through matrix-vector products.

    Might be useful if for the eigendecomposition of a matrix :math:`A = UU^\top`,
    where :math:`U` is a known matrix of size :math:`(n,m)`;
    computing :math:`U(U^\top x)` costs :math:`O(nm)` operations whereas
    computing :math:`A` costs :math:`O(n^2m)`.

    See :meth:`fastpca.eigsh` for a description of the other arguments.

    Arguments:
        f (function): Applies the linear operator to eigendecompose,
            :math:`f(x) = Ax`.
        dtype (torch.dtype): The type used in ``f``.
        device (torch.device): The device where ``f`` is allocated.
        m (int): The dimensionality of the domain of ``f``.
            Would be ``A.shape[1]`` if :math:`f(x) = Ax`.
    """
    if L is None:
        L = k+2
    assert k >= 0
    assert k <= n
    assert L >= k
    assert n_iter >= 0

    def orthogonalize(A):
        (Q, _) = torch.qr(A)
        return Q

    def nystrom(Q, anorm):
        r"""
        Use the Nystrom method to obtain approximations to the
        eigenvalues and eigenvectors of A (shifting A on the subspace
        spanned by the columns of Q in order to make the shifted A be
        positive definite).
        """
        
        def svd_thin_matrix(A):
            r"""
            Efficient implementation of SVD on [N x D] matrix, D >> N.
            """
            (e, V) = torch.symeig(A @ A.t(), eigenvectors=True)
            
            Sigma = torch.sqrt(e)
            SigInv = 1/Sigma 
            SigInv[torch.isnan(SigInv)] = 0
            U = A.t() @ (V * SigInv)
            
            return U, Sigma, V
            
        anorm = .1e-6 * anorm * math.sqrt(1. * n)
        E = f(Q) + anorm * Q
        R = Q.t() @ E
        R = (R + R.t()) / 2
        R = torch.potrf(R, upper=False) # Cholesky
        (tmp, _) = torch.gesv(E.t(), R) # Solve
        V, d, _ = svd_thin_matrix(tmp)
        d = d * d - anorm
        return d, V

    Q = 2*(torch.rand((n, L), device=device, dtype=dtype)-0.5)
    for _ in range(max(0, n_iter-1)):
        Q = orthogonalize(f(Q))
    oldQ = Q
    Q = f(Q)
    anorm = torch.max(torch.norm(Q, dim=0)/torch.norm(oldQ, dim=0))
    Q = orthogonalize(Q)

    d, V = nystrom(Q, anorm)

    # Retain only the entries with the k greatest absolute values
    (_, idx) = torch.abs(d).sort()
    idx = idx[(L-k):]
    return abs(d[idx]), V[:, idx]
//  [python]
def eigsh_func(f, dtype, device, n, k=6, n_iter=4, L=None):
    r"""Top-k Eigendecomposition of PSD linear operator :math:`f(x) = Ax`.

    Provides a way to use :meth:`fastpca.eigsh` when the matrix
    to eigendecompose is only accessible through matrix-vector products.

    Might be useful if for the eigendecomposition of a matrix :math:`A = UU^\top`,
    where :math:`U` is a known matrix of size :math:`(n,m)`;
    computing :math:`U(U^\top x)` costs :math:`O(nm)` operations whereas
    computing :math:`A` costs :math:`O(n^2m)`.

    See :meth:`fastpca.eigsh` for a description of the other arguments.

    Arguments:
        f (function): Applies the linear operator to eigendecompose,
            :math:`f(x) = Ax`.
        dtype (torch.dtype): The type used in ``f``.
        device (torch.device): The device where ``f`` is allocated.
        m (int): The dimensionality of the domain of ``f``.
            Would be ``A.shape[1]`` if :math:`f(x) = Ax`.
    """
    if L is None:
        L = k+2
    assert k >= 0
    assert k <= n
    assert L >= k
    assert n_iter >= 0

    def orthogonalize(A):
        (Q, _) = torch.linalg.qr(A)
        return Q

    def nystrom(Q, anorm):
        r"""
        Use the Nystrom method to obtain approximations to the
        eigenvalues and eigenvectors of A (shifting A on the subspace
        spanned by the columns of Q in order to make the shifted A be
        positive definite).
        """
        
        def svd_thin_matrix(A):
            r"""
            Efficient implementation of SVD on [N x D] matrix, D >> N.
            """
            (e, V) = torch.symeig(A @ A.t(), eigenvectors=True)
            
            Sigma = torch.sqrt(e)
            SigInv = 1/Sigma 
            SigInv[torch.isnan(SigInv)] = 0
            U = A.t() @ (V * SigInv)
            
            return U, Sigma, V
            
        anorm = .1e-6 * anorm * math.sqrt(1. * n)
        E = f(Q) + anorm * Q
        R = Q.t() @ E
        R = (R + R.t()) / 2
        R = torch.potrf(R, upper=False) # Cholesky
        (tmp, _) = torch.gesv(E.t(), R) # Solve
        V, d, _ = svd_thin_matrix(tmp)
        d = d * d - anorm
        return d, V

    Q = 2*(torch.rand((n, L), device=device, dtype=dtype)-0.5)
    for _ in range(max(0, n_iter-1)):
        Q = orthogonalize(f(Q))
    oldQ = Q
    Q = f(Q)
    anorm = torch.max(torch.norm(Q, dim=0)/torch.norm(oldQ, dim=0))
    Q = orthogonalize(Q)

    d, V = nystrom(Q, anorm)

    # Retain only the entries with the k greatest absolute values
    (_, idx) = torch.abs(d).sort()
    idx = idx[(L-k):]
    return abs(d[idx]), V[:, idx]
//  [python]
def __init__(self, c): 
        super(Invertible1x1Conv, self).__init__() 
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0, 
                                    bias=False) 
 
        # Sample a random orthonormal matrix to initialize weights 
        W = torch.qr(torch.FloatTensor(c, c).normal_())[0] 
 
        # Ensure determinant is 1.0 not -1.0 
        if torch.det(W) < 0: 
            W[:,0] = -1*W[:,0] 
        W = W.view(c, c, 1) 
        self.conv.weight.data = W
//  [python]
def __init__(self, c): 
        super(Invertible1x1Conv, self).__init__() 
        self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0, 
                                    bias=False) 
 
        # Sample a random orthonormal matrix to initialize weights 
        W = torch.linalg.qr(torch.FloatTensor(c, c).normal_())[0] 
 
        # Ensure determinant is 1.0 not -1.0 
        if torch.det(W) < 0: 
            W[:,0] = -1*W[:,0] 
        W = W.view(c, c, 1) 
        self.conv.weight.data = W
//  [python]
def orthogonal_matrix_chunk(cols, qr_uniform_q = False, device = None):
    # MARK: change cols (float 64.0) to int(cols)
    unstructured_block = torch.randn((int(cols), int(cols)), device = device)
    q, r = torch.qr(unstructured_block.cpu(), some = True)
    q, r = map(lambda t: t.to(device), (q, r))

    # proposed by @Parskatt
    # to make sure Q is uniform https://arxiv.org/pdf/math-ph/0609050.pdf
    if qr_uniform_q:
        d = torch.diag(r, 0)
        q *= d.sign()
    return q.t()
//  [python]
def orthogonal_matrix_chunk(cols, qr_uniform_q = False, device = None):
    # MARK: change cols (float 64.0) to int(cols)
    unstructured_block = torch.randn((int(cols), int(cols)), device = device)
    q, r = torch.linalg.qr(unstructured_block.cpu())
    q, r = map(lambda t: t.to(device), (q, r))

    # proposed by @Parskatt
    # to make sure Q is uniform https://arxiv.org/pdf/math-ph/0609050.pdf
    if qr_uniform_q:
        d = torch.diag(r, 0)
        q *= d.sign()
    return q.t()
//  [python]
def __init__(self, x_shape, num_u_channels=0):
        assert len(x_shape) == 1 or len(x_shape) == 3
        super().__init__(x_shape, x_shape)

        num_channels = x_shape[0]
        self.weight_shape = [num_channels, num_channels]

        self.conv_weights_shape = self.weight_shape + [1 for _ in x_shape[1:]]
        self.num_non_channel_elements = np.prod(x_shape[1:])

        self.weights_init = torch.qr(torch.randn(*self.weight_shape))[0]

        self.num_u_channels = num_u_channels
        if num_u_channels > 0:
            self.u_weights = nn.Parameter(torch.zeros(num_channels, num_u_channels))
            self.u_conv_weights_shape = [num_channels, num_u_channels] + [1 for _ in x_shape[1:]]
//  [python]
def __init__(self, x_shape, num_u_channels=0):
        assert len(x_shape) == 1 or len(x_shape) == 3
        super().__init__(x_shape, x_shape)

        num_channels = x_shape[0]
        self.weight_shape = [num_channels, num_channels]

        self.conv_weights_shape = self.weight_shape + [1 for _ in x_shape[1:]]
        self.num_non_channel_elements = np.prod(x_shape[1:])

        self.weights_init = torch.linalg.qr(torch.randn(*self.weight_shape))[0]

        self.num_u_channels = num_u_channels
        if num_u_channels > 0:
            self.u_weights = nn.Parameter(torch.zeros(num_channels, num_u_channels))
            self.u_conv_weights_shape = [num_channels, num_u_channels] + [1 for _ in x_shape[1:]]
//  [python]
def orthogonalize_weights(tensor, act):
    r"""Fills the input `Tensor` with a (semi) orthogonal matrix, as
    described in `Exact solutions to the nonlinear dynamics of learning in deep
    linear neural networks` - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the
    trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional `torch.Tensor`, where :math:`n \geq 2`
        gain: optional scaling factor

    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.orthogonal_(w)
    """
    tensor = tensor.clone().detach() # @mst: avoid modifying the original tensor
    gain = init.calculate_gain(act)

    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor.numel() // rows
    # flattened = tensor.new(rows, cols).normal_(0, 1)
    flattened = tensor.view(rows, cols) # @mst: do NOT reinit the tensor

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.qr(flattened) # q: (rows, cols), r: (cols, cols) if rows > cols.

    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0) # d: (cols)
    ph = d.sign()
    q *= ph # (rows, cols) * (cols) => (rows, cols)

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def orthogonalize_weights(tensor, act):
    r"""Fills the input `Tensor` with a (semi) orthogonal matrix, as
    described in `Exact solutions to the nonlinear dynamics of learning in deep
    linear neural networks` - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the
    trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional `torch.Tensor`, where :math:`n \geq 2`
        gain: optional scaling factor

    Examples:
        >>> w = torch.empty(3, 5)
        >>> nn.init.orthogonal_(w)
    """
    tensor = tensor.clone().detach() # @mst: avoid modifying the original tensor
    gain = init.calculate_gain(act)

    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor.numel() // rows
    # flattened = tensor.new(rows, cols).normal_(0, 1)
    flattened = tensor.view(rows, cols) # @mst: do NOT reinit the tensor

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.linalg.qr(flattened) # q: (rows, cols), r: (cols, cols) if rows > cols.

    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0) # d: (cols)
    ph = d.sign()
    q *= ph # (rows, cols) * (cols) => (rows, cols)

    if rows < cols:
        q.t_()

    with torch.no_grad():
        tensor.view_as(q).copy_(q)
        tensor.mul_(gain)
    return tensor
//  [python]
def __init__(self, in_channel, out_channel=None):
        super().__init__()
        if out_channel is None:
            out_channel = in_channel
        weight = torch.randn(in_channel, out_channel)
        q, _ = torch.qr(weight)
        weight = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(weight)
//  [python]
def __init__(self, in_channel, out_channel=None):
        super().__init__()
        if out_channel is None:
            out_channel = in_channel
        weight = torch.randn(in_channel, out_channel)
        q, _ = torch.linalg.qr(weight)
        weight = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(weight)
//  [python]
def __init__(self, channels, num_splits=4, no_jacobian=False, **kwargs):  # pylint: disable=unused-argument
        super().__init__()
        assert num_splits % 2 == 0
        self.channels = channels
        self.num_splits = num_splits
        self.no_jacobian = no_jacobian
        self.weight_inv = None

        if Version(torch.__version__) < Version("1.9"):
            w_init = torch.qr(torch.FloatTensor(self.num_splits, self.num_splits).normal_())[0]
        else:
            w_init = torch.linalg.qr(torch.FloatTensor(self.num_splits, self.num_splits).normal_(), "complete")[0]

        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.weight = nn.Parameter(w_init)
//  [python]
def __init__(self, channels, num_splits=4, no_jacobian=False, **kwargs):  # pylint: disable=unused-argument
        super().__init__()
        assert num_splits % 2 == 0
        self.channels = channels
        self.num_splits = num_splits
        self.no_jacobian = no_jacobian
        self.weight_inv = None

        w_init = torch.linalg.qr(torch.FloatTensor(self.num_splits, self.num_splits).normal_(), "complete")[0]

        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.weight = nn.Parameter(w_init)
//  [python]
def orthogonal(tensor, gain=1):
    """Fills the input Tensor or Variable with a (semi) orthogonal matrix, as
    described in "Exact solutions to the nonlinear dynamics of learning in deep
    linear neural networks" - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the
    trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional torch.Tensor or autograd.Variable, where n >= 2
        gain: optional scaling factor

    Examples:
        >>> w = torch.Tensor(3, 5)
        >>> nn.init.orthogonal(w)
    """
    if isinstance(tensor, Variable):
        orthogonal(tensor.data, gain=gain)
        return tensor

    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor[0].numel()
    flattened = torch.Tensor(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph.expand_as(q)

    if rows < cols:
        q.t_()

    tensor.view_as(q).copy_(q)
    tensor.mul_(gain)
    return tensor
//  [python]
def orthogonal(tensor, gain=1):
    """Fills the input Tensor or Variable with a (semi) orthogonal matrix, as
    described in "Exact solutions to the nonlinear dynamics of learning in deep
    linear neural networks" - Saxe, A. et al. (2013). The input tensor must have
    at least 2 dimensions, and for tensors with more than 2 dimensions the
    trailing dimensions are flattened.

    Args:
        tensor: an n-dimensional torch.Tensor or autograd.Variable, where n >= 2
        gain: optional scaling factor

    Examples:
        >>> w = torch.Tensor(3, 5)
        >>> nn.init.orthogonal(w)
    """
    if isinstance(tensor, Variable):
        orthogonal(tensor.data, gain=gain)
        return tensor

    if tensor.ndimension() < 2:
        raise ValueError("Only tensors with 2 or more dimensions are supported")

    rows = tensor.size(0)
    cols = tensor[0].numel()
    flattened = torch.Tensor(rows, cols).normal_(0, 1)

    if rows < cols:
        flattened.t_()

    # Compute the qr factorization
    q, r = torch.linalg.qr(flattened)
    # Make Q uniform according to https://arxiv.org/pdf/math-ph/0609050.pdf
    d = torch.diag(r, 0)
    ph = d.sign()
    q *= ph.expand_as(q)

    if rows < cols:
        q.t_()

    tensor.view_as(q).copy_(q)
    tensor.mul_(gain)
    return tensor
//  [python]
def test_aten_fallback(self):
        class ModelWithAtenNotONNXOp(nn.Module):
            def forward(self, x, y):
                abcd = x + y
                defg = torch.qr(abcd)
                return defg

        x = torch.rand(3, 4)
        y = torch.rand(3, 4)
        f = io.BytesIO()
        torch.onnx.export_to_pretty_string(
            ModelWithAtenNotONNXOp(), (x, y), f,
            add_node_names=False,
            do_constant_folding=False,
            operator_export_type=OperatorExportTypes.ONNX_ATEN_FALLBACK)
//  [python]
def test_aten_fallback(self):
        class ModelWithAtenNotONNXOp(nn.Module):
            def forward(self, x, y):
                abcd = x + y
                defg = torch.linalg.qr(abcd)
                return defg

        x = torch.rand(3, 4)
        y = torch.rand(3, 4)
        f = io.BytesIO()
        torch.onnx.export_to_pretty_string(
            ModelWithAtenNotONNXOp(), (x, y), f,
            add_node_names=False,
            do_constant_folding=False,
            operator_export_type=OperatorExportTypes.ONNX_ATEN_FALLBACK)
//  [python]
def __init__(self, feat_dim, flag_detjac=False):
        super(Invertible1x1ConvWaveGlow, self).__init__()

        torch.manual_seed(100)
        with torch.no_grad():
            W = torch.qr(torch.FloatTensor(feat_dim, feat_dim).normal_())[0]
            
            # Ensure determinant is 1.0 not -1.0
            if torch.det(W) < 0:
                W[:,0] = -1*W[:,0]
                
            # not necessary
            W = W.transpose(0, 1)
        self.weight = torch_nn.Parameter(W)
        self.weight_inv = torch_nn.Parameter(W.clone())
        self.weight_inv_flag = False
        self.flag_detjac = flag_detjac
        return
//  [python]
def __init__(self, feat_dim, flag_detjac=False):
        super(Invertible1x1ConvWaveGlow, self).__init__()

        torch.manual_seed(100)
        with torch.no_grad():
            W = torch.linalg.qr(torch.FloatTensor(feat_dim, feat_dim).normal_())[0]
            
            # Ensure determinant is 1.0 not -1.0
            if torch.det(W) < 0:
                W[:,0] = -1*W[:,0]
                
            # not necessary
            W = W.transpose(0, 1)
        self.weight = torch_nn.Parameter(W)
        self.weight_inv = torch_nn.Parameter(W.clone())
        self.weight_inv_flag = False
        self.flag_detjac = flag_detjac
        return
//  [python]
def test_qr(self):
        for i in range(50):
            tn = np.random.randn(3, 3).astype('float32')
            while np.allclose(np.linalg.det(tn), 0):
                tn = np.random.randn((3, 3)).astype('float32')
            x = jt.array(tn)
            # x = x.reindex([2, 2, x.shape[0], x.shape[1]], ["i2", "i3"])
            t_x = torch.from_numpy(tn)
            t_x = Variable(t_x, requires_grad=True)
            jq, jr = jt.linalg.qr(x)
            tq, tr = torch.qr(t_x)
            try:
                assert np.allclose(jq.data, tq.detach().numpy(), rtol=1e-4, atol=1e-6)
                assert np.allclose(jr.data, tr.detach().numpy(), rtol=1e-4, atol=1e-6)
            except AssertionError:
                print("ours' qr results:")
                print(jq)
                print(jr)
                print("pytorch's qr results:")
                print(tq)
                print(tr)
            gq = jt.grad(jq, x).data
            gr = jt.grad(jr, x).data
            tgq = torch.autograd.grad(tq, t_x, torch.ones_like(tq), retain_graph=True)
            tgr = torch.autograd.grad(tr, t_x, torch.ones_like(tr), retain_graph=True)
            try:
                assert np.allclose(gq, tgq[0].numpy(), rtol=1e-4, atol=1e-6)
                assert np.allclose(gr, tgr[0].numpy(), rtol=1e-4, atol=1e-6)
            except AssertionError:
                print("ours' qr grad results:")
                print(gq)
                print(gr)
                print("pytorch's qr grad result")
                print(tgq[0])
                print(tgr[0])
//  [python]
def test_qr(self):
        for i in range(50):
            tn = np.random.randn(3, 3).astype('float32')
            while np.allclose(np.linalg.det(tn), 0):
                tn = np.random.randn((3, 3)).astype('float32')
            x = jt.array(tn)
            # x = x.reindex([2, 2, x.shape[0], x.shape[1]], ["i2", "i3"])
            t_x = torch.from_numpy(tn)
            t_x = Variable(t_x, requires_grad=True)
            jq, jr = jt.linalg.qr(x)
            tq, tr = torch.linalg.qr(t_x)
            try:
                assert np.allclose(jq.data, tq.detach().numpy(), rtol=1e-4, atol=1e-6)
                assert np.allclose(jr.data, tr.detach().numpy(), rtol=1e-4, atol=1e-6)
            except AssertionError:
                print("ours' qr results:")
                print(jq)
                print(jr)
                print("pytorch's qr results:")
                print(tq)
                print(tr)
            gq = jt.grad(jq, x).data
            gr = jt.grad(jr, x).data
            tgq = torch.autograd.grad(tq, t_x, torch.ones_like(tq), retain_graph=True)
            tgr = torch.autograd.grad(tr, t_x, torch.ones_like(tr), retain_graph=True)
            try:
                assert np.allclose(gq, tgq[0].numpy(), rtol=1e-4, atol=1e-6)
                assert np.allclose(gr, tgr[0].numpy(), rtol=1e-4, atol=1e-6)
            except AssertionError:
                print("ours' qr grad results:")
                print(gq)
                print(gr)
                print("pytorch's qr grad result")
                print(tgq[0])
                print(tgr[0])
//  [python]
def __init__(self, c):
		super(Invertible1x1Conv, self).__init__()
		self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0,
									bias=False)

		# Sample a random orthonormal matrix to initialize weights
		W = torch.qr(torch.FloatTensor(c, c).normal_())[0]

		# Ensure determinant is 1.0 not -1.0
		if torch.det(W) < 0:
			W[:,0] = -1*W[:,0]
		W = W.view(c, c, 1)
		self.conv.weight.data = W
//  [python]
def __init__(self, c):
		super(Invertible1x1Conv, self).__init__()
		self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0,
									bias=False)

		# Sample a random orthonormal matrix to initialize weights
		W = torch.linalg.qr(torch.FloatTensor(c, c).normal_())[0]

		# Ensure determinant is 1.0 not -1.0
		if torch.det(W) < 0:
			W[:,0] = -1*W[:,0]
		W = W.view(c, c, 1)
		self.conv.weight.data = W
//  [python]
def forward(self, feature, mask = None):
        p = self.patch_size
        # transformer
        x = rearrange(feature, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = p[0], p2 = p[1])
        # print(x.shape)
        x = self.patch_to_embedding(x)

        b, n, _ = x.shape

        mq,mr = torch.qr(self.motion_token+1e-9)
        motion_tokens = mq.permute(0,2,1).repeat(b,1,1).contiguous()
        if self.pos_embedding_type in ["sine","sine-full"] :
            x += self.pos_embedding[:, :n]
            x = torch.cat((motion_tokens.detach(), x), dim=1)
        elif self.pos_embedding_type == 'none':
            x = torch.cat((motion_tokens.detach(), x), dim=1)
        else:
            x = torch.cat((motion_tokens.detach(), x), dim=1)
            x += self.pos_embedding[:, :(n + self.num_motions)]
        x = self.dropout(x)

        x = self.transformer(x, mask,self.pos_embedding)
        x = self.to_motion_token(x[:, 0:self.num_motions])
        weights = self.mlp_head(x)
        x = x * weights
        return x
//  [python]
def forward(self, feature, mask = None):
        p = self.patch_size
        # transformer
        x = rearrange(feature, 'b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = p[0], p2 = p[1])
        # print(x.shape)
        x = self.patch_to_embedding(x)

        b, n, _ = x.shape

        mq,mr = torch.linalg.qr(self.motion_token+1e-9)
        motion_tokens = mq.permute(0,2,1).repeat(b,1,1).contiguous()
        if self.pos_embedding_type in ["sine","sine-full"] :
            x += self.pos_embedding[:, :n]
            x = torch.cat((motion_tokens.detach(), x), dim=1)
        elif self.pos_embedding_type == 'none':
            x = torch.cat((motion_tokens.detach(), x), dim=1)
        else:
            x = torch.cat((motion_tokens.detach(), x), dim=1)
            x += self.pos_embedding[:, :(n + self.num_motions)]
        x = self.dropout(x)

        x = self.transformer(x, mask,self.pos_embedding)
        x = self.to_motion_token(x[:, 0:self.num_motions])
        weights = self.mlp_head(x)
        x = x * weights
        return x
//  [python]
def __init__(self, in_channel):
        super().__init__()
        q, _ = torch.qr(torch.randn(in_channel, in_channel))
        # making it 1x1 conv: conv2d(in_channels=in_channel, out_channels=in_channel, kernel_size=1, stride=1)
        w = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(w)
//  [python]
def __init__(self, in_channel):
        super().__init__()
        q, _ = torch.linalg.qr(torch.randn(in_channel, in_channel))
        # making it 1x1 conv: conv2d(in_channels=in_channel, out_channels=in_channel, kernel_size=1, stride=1)
        w = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(w)