jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def xnystrace(A: lo.LinearOperator, m: int=80, *, _gen_data_on_device=None):
    n, n1 = A.shape
    assert n == n1
    if n <= m:
        return exact_trace(A)
    m = m // 2
    def normalize_columns(M):
        return M / torch.linalg.vector_norm(M, dim=0)

    def diag_of_AB(A, B):
        return torch.sum(A * B, dim=0)

    Z = _sample_dist('gaussian', (n, m), _gen_data_on_device, A.device)
    Omega = math.sqrt(n) * normalize_columns(Z)
    Y = A @ Omega

    nu = torch.finfo(Omega.dtype).eps * torch.linalg.vector_norm(Y) / math.sqrt(n)
    Y = Y + nu * Omega


    Q, R = torch.linalg.qr(Y)
    H = Omega.T @ Y
    C = torch.linalg.cholesky((H + H.T) / 2, upper=True)

    B = torch.linalg.solve_triangular(C, R, upper=True, left=False) # Double check this

    QQ, RR = torch.linalg.qr(Omega)
    WW = QQ.T @ Omega
    SS = normalize_columns(torch.linalg.inv(RR).T)

    scale = (n - m + 1) / (n - torch.linalg.vector_norm(WW, dim=0)**2 +
                           torch.abs(diag_of_AB(SS, WW) * torch.linalg.vector_norm(SS, dim=0))**2)
    W = Q.T @ Omega
    S =  torch.linalg.solve_triangular(C.T, B, upper=False, left=False) * torch.diag(torch.linalg.inv(H))**(-0.5) # Double check triangular solve
    dSW = diag_of_AB(S, W)
    ests = torch.linalg.norm(B, 'fro')**2 - torch.linalg.vector_norm(S, dim=0)**2 + torch.abs(dSW)**2 * scale - nu * n
    return torch.mean(ests), torch.std(ests) / math.sqrt(m)
//  [python]
def random_orthogonal_matrix(rows: int, cols: int) -> torch.Tensor:
    """Returns a random orthogonal matrix of shape (rows, cols)."""
    q, _ = torch.qr(torch.randn(rows, cols))
    return q
//  [python]
def random_orthogonal_matrix(rows: int, cols: int) -> torch.Tensor:
    """Returns a random orthogonal matrix of shape (rows, cols)."""
    q, _ = torch.linalg.qr(torch.randn(rows, cols))
    return q
//  [python]
def make_orthogonal(tensor):
    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = torch.reshape(tensor, (rows, cols))

    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_()
        
    return q
//  [python]
def make_orthogonal(tensor):
    rows = tensor.size(0)
    cols = tensor.numel() // rows
    flattened = torch.reshape(tensor, (rows, cols))

    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_()
        
    return q
//  [python]
np.random.seed(1)
torch.random.manual_seed(1)

N = 100
gamma = 0.001

U, _ = torch.qr(torch.randn((N, N)))
//  [python]
np.random.seed(1)
torch.random.manual_seed(1)

N = 100
gamma = 0.001

U, _ = torch.linalg.qr(torch.randn((N, N)))
//  [python]
def params2orb(params, coeffs, with_penalty):
        orb, _ = torch.qr(params)
        if with_penalty is None:
            return orb
        else:
            # QR decomposition's solution is not unique in a way that every column
            # can be multiplied by -1 and it still a solution
            # So, to remove the non-uniqueness, we will make the sign of the sum
            # positive.
            s1 = torch.sign(orb.sum(dim=-2, keepdim=True))  # (*BD, 1, norb)
            s2 = torch.sign(params.sum(dim=-2, keepdim=True))
            penalty = torch.mean((orb * s1 - params * s2) ** 2) * with_penalty
            return orb, penalty
//  [python]
def params2orb(params, coeffs, with_penalty):
        orb, _ = torch.linalg.qr(params)
        if with_penalty is None:
            return orb
        else:
            # QR decomposition's solution is not unique in a way that every column
            # can be multiplied by -1 and it still a solution
            # So, to remove the non-uniqueness, we will make the sign of the sum
            # positive.
            s1 = torch.sign(orb.sum(dim=-2, keepdim=True))  # (*BD, 1, norb)
            s2 = torch.sign(params.sum(dim=-2, keepdim=True))
            penalty = torch.mean((orb * s1 - params * s2) ** 2) * with_penalty
            return orb, penalty
//  [python]
def test_ormqr(self, device, dtype):

        def run_test(batch, m, n, fortran_contiguous):
            A = make_tensor((*batch, m, n), dtype=dtype, device=device)
            reflectors, tau = torch.geqrf(A)
            if not fortran_contiguous:
                self.assertTrue(reflectors.transpose(-2, -1).is_contiguous())
                reflectors = reflectors.contiguous()

            # Q is of size m x m
            Q, _ = torch.qr(A, some=True)
            C_right = make_tensor((*batch, m, n), dtype=dtype, device=device)
            C_left = make_tensor((*batch, n, m), dtype=dtype, device=device)

            expected = Q @ C_right
            actual = torch.ormqr(reflectors, tau, C_right, left=True, transpose=False)
            self.assertEqual(expected, actual)

            expected = C_left @ Q
            actual = torch.ormqr(reflectors, tau, C_left, left=False, transpose=False)
            self.assertEqual(expected, actual)

            expected = Q.transpose(-2, -1).conj() @ C_right
            actual = torch.ormqr(reflectors, tau, C_right, left=True, transpose=True)
            self.assertEqual(expected, actual)

            expected = C_left @ Q.transpose(-2, -1).conj()
            actual = torch.ormqr(reflectors, tau, C_left, left=False, transpose=True)
            self.assertEqual(expected, actual)

            # if tau is all zeros then the implicit matrix Q is the identity matrix
            # so the actual result should be C_right in this case
            zero_tau = torch.zeros_like(tau)
            actual = torch.ormqr(reflectors, zero_tau, C_right, left=True, transpose=False)
            self.assertEqual(C_right, actual)

        batches = [(), (0, ), (2, ), (2, 1)]
        ns = [5, 2, 0]
        for batch, (m, n), fortran_contiguous in product(batches, product(ns, ns), [True, False]):
            run_test(batch, m, n, fortran_contiguous)
//  [python]
def test_ormqr(self, device, dtype):

        def run_test(batch, m, n, fortran_contiguous):
            A = make_tensor((*batch, m, n), dtype=dtype, device=device)
            reflectors, tau = torch.geqrf(A)
            if not fortran_contiguous:
                self.assertTrue(reflectors.transpose(-2, -1).is_contiguous())
                reflectors = reflectors.contiguous()

            # Q is of size m x m
            Q, _ = torch.linalg.qr(A, mode='complete')
            C_right = make_tensor((*batch, m, n), dtype=dtype, device=device)
            C_left = make_tensor((*batch, n, m), dtype=dtype, device=device)

            expected = Q @ C_right
            actual = torch.ormqr(reflectors, tau, C_right, left=True, transpose=False)
            self.assertEqual(expected, actual)

            expected = C_left @ Q
            actual = torch.ormqr(reflectors, tau, C_left, left=False, transpose=False)
            self.assertEqual(expected, actual)

            expected = Q.transpose(-2, -1).conj() @ C_right
            actual = torch.ormqr(reflectors, tau, C_right, left=True, transpose=True)
            self.assertEqual(expected, actual)

            expected = C_left @ Q.transpose(-2, -1).conj()
            actual = torch.ormqr(reflectors, tau, C_left, left=False, transpose=True)
            self.assertEqual(expected, actual)

            # if tau is all zeros then the implicit matrix Q is the identity matrix
            # so the actual result should be C_right in this case
            zero_tau = torch.zeros_like(tau)
            actual = torch.ormqr(reflectors, zero_tau, C_right, left=True, transpose=False)
            self.assertEqual(C_right, actual)

        batches = [(), (0, ), (2, ), (2, 1)]
        ns = [5, 2, 0]
        for batch, (m, n), fortran_contiguous in product(batches, product(ns, ns), [True, False]):
            run_test(batch, m, n, fortran_contiguous)
//  [python]
def gaussian_orthogonal_random_matrix(nrows, ncols, scaling=0, device=None, dtype=None):
    factory_kwargs = {'device': device, 'dtype': dtype}
    nblocks = int(math.ceil(nrows / ncols))
    # TD [2021-10-28]: Sometimes QR fails on CUDA
    unstructured_blocks = torch.randn((nblocks, ncols, ncols), device='cpu')
    q, r = torch.qr(unstructured_blocks)
    # To make sure Q is uniform from the Haar distribution https://arxiv.org/pdf/math-ph/0609050.pdf
    q *= rearrange(torch.diagonal(r, dim1=-2, dim2=-1).sign(), 'b c -> b 1 c')
    q = q.to(**factory_kwargs)
    # TD [2021-10-28] Idk why the transpose is necessary. I suspect it isn't.
    # https://github.com/google-research/google-research/blob/ea313c6e96acce6c863de41615c6cf4079b8ca94/performer/fast_attention/jax/fast_attention.py#L362
    q = rearrange(q, 'b c c1 -> b c1 c')
    g_ortho = rearrange(q, 'b c1 c -> (b c1) c')[:nrows]

    if scaling == 0:
        multiplier = torch.randn((nrows, ncols), **factory_kwargs).norm(dim=1)
        return rearrange(multiplier, 'r -> r 1') * g_ortho
    elif scaling == 1:
        return math.sqrt(ncols) * g_ortho
    else:
        raise ValueError(f'Invalid scaling {scaling}')
//  [python]
def gaussian_orthogonal_random_matrix(nrows, ncols, scaling=0, device=None, dtype=None):
    factory_kwargs = {'device': device, 'dtype': dtype}
    nblocks = int(math.ceil(nrows / ncols))
    # TD [2021-10-28]: Sometimes QR fails on CUDA
    unstructured_blocks = torch.randn((nblocks, ncols, ncols), device='cpu')
    q, r = torch.linalg.qr(unstructured_blocks)
    # To make sure Q is uniform from the Haar distribution https://arxiv.org/pdf/math-ph/0609050.pdf
    q *= rearrange(torch.diagonal(r, dim1=-2, dim2=-1).sign(), 'b c -> b 1 c')
    q = q.to(**factory_kwargs)
    # TD [2021-10-28] Idk why the transpose is necessary. I suspect it isn't.
    # https://github.com/google-research/google-research/blob/ea313c6e96acce6c863de41615c6cf4079b8ca94/performer/fast_attention/jax/fast_attention.py#L362
    q = rearrange(q, 'b c c1 -> b c1 c')
    g_ortho = rearrange(q, 'b c1 c -> (b c1) c')[:nrows]

    if scaling == 0:
        multiplier = torch.randn((nrows, ncols), **factory_kwargs).norm(dim=1)
        return rearrange(multiplier, 'r -> r 1') * g_ortho
    elif scaling == 1:
        return math.sqrt(ncols) * g_ortho
    else:
        raise ValueError(f'Invalid scaling {scaling}')
//  [python]
def params2orb(params: torch.Tensor,
                   coeffs: torch.Tensor,
                   with_penalty: float = 0.0) -> List[torch.Tensor]:
        """
        Convert the parameters & coefficients to the orthogonal orbitals.
        ``params`` is the tensor to be optimized in variational method, while
        ``coeffs`` is a tensor that is needed to get the orbital, but it is not
        optimized in the variational method.

        Parameters
        ----------
        params: torch.Tensor
            The free parameters to be optimized.
        coeffs: torch.Tensor
            The coefficients to get the orthogonal orbitals.
        with_penalty: float (default 0.0)
            If not 0.0, return the penalty term for the free parameters.

        Returns
        -------
        orb: torch.Tensor
            The orthogonal orbitals.
        penalty: torch.Tensor
            The penalty term for the free parameters. If ``with_penalty`` is 0.0,
            this is not returned.

        """
        q, r = torch.qr(params)
        orb = q
        if with_penalty == 0.0:
            return [orb]
        else:
            # QR decomposition's solution is not unique in a way that every column
            # can be multiplied by -1 and it still a solution
            # So, to remove the non-uniqueness, we will make the sign of the sum
            # positive.
            s1 = torch.sign(orb.sum(dim=-2, keepdim=True))  # (*BD, 1, norb)
            s2 = torch.sign(params.sum(dim=-2, keepdim=True))
            penalty = torch.mean((orb * s1 - params * s2)**2) * with_penalty
            return [orb, penalty]
//  [python]
def params2orb(params: torch.Tensor,
                   coeffs: torch.Tensor,
                   with_penalty: float = 0.0) -> List[torch.Tensor]:
        """
        Convert the parameters & coefficients to the orthogonal orbitals.
        ``params`` is the tensor to be optimized in variational method, while
        ``coeffs`` is a tensor that is needed to get the orbital, but it is not
        optimized in the variational method.

        Parameters
        ----------
        params: torch.Tensor
            The free parameters to be optimized.
        coeffs: torch.Tensor
            The coefficients to get the orthogonal orbitals.
        with_penalty: float (default 0.0)
            If not 0.0, return the penalty term for the free parameters.

        Returns
        -------
        orb: torch.Tensor
            The orthogonal orbitals.
        penalty: torch.Tensor
            The penalty term for the free parameters. If ``with_penalty`` is 0.0,
            this is not returned.

        """
        orb, _ = torch.linalg.qr(params)
        if with_penalty == 0.0:
            return [orb]
        else:
            # QR decomposition's solution is not unique in a way that every column
            # can be multiplied by -1 and it still a solution
            # So, to remove the non-uniqueness, we will make the sign of the sum
            # positive.
            s1 = torch.sign(orb.sum(dim=-2, keepdim=True))  # (*BD, 1, norb)
            s2 = torch.sign(params.sum(dim=-2, keepdim=True))
            penalty = torch.mean((orb * s1 - params * s2)**2) * with_penalty
            return [orb, penalty]
//  [python]
def __init__(self, channels, n_split=4, no_jacobian=False, **kwargs):
        super().__init__()
        assert n_split % 2 == 0
        self.channels = channels
        self.n_split = n_split
        self.no_jacobian = no_jacobian

        w_init = torch.qr(
            torch.FloatTensor(self.n_split, self.n_split).normal_()
        )[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, n_split=4, no_jacobian=False, **kwargs):
        super().__init__()
        assert n_split % 2 == 0
        self.channels = channels
        self.n_split = n_split
        self.no_jacobian = no_jacobian

        w_init = torch.linalg.qr(
            torch.FloatTensor(self.n_split, self.n_split).normal_()
        )[0]
        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.weight = nn.Parameter(w_init)
//  [python]
def random_orthogonal_matrix(size, device):
    """
    Generate a random orthogonal matrix of the specified size.
    First, we generate a random matrix with entries from a standard distribution.
    Then, we use QR decomposition to obtain an orthogonal matrix.
    Finally, we multiply by a diagonal matrix with diag r to adjust the signs.
    
    Args:
    size (int): The size of the matrix (size x size).
    
    Returns:
    torch.Tensor: An orthogonal matrix of the specified size.
    """
    torch.cuda.empty_cache()
    random_matrix = torch.randn(size, size, dtype=torch.float64).to(device)
    q, r = torch.qr(random_matrix)
    q *= torch.sign(torch.diag(r)).unsqueeze(0)
    return q
//  [python]
def random_orthogonal_matrix(size, device):
    """
    Generate a random orthogonal matrix of the specified size.
    First, we generate a random matrix with entries from a standard distribution.
    Then, we use QR decomposition to obtain an orthogonal matrix.
    Finally, we multiply by a diagonal matrix with diag r to adjust the signs.
    
    Args:
    size (int): The size of the matrix (size x size).
    
    Returns:
    torch.Tensor: An orthogonal matrix of the specified size.
    """
    torch.cuda.empty_cache()
    random_matrix = torch.randn(size, size, dtype=torch.float64).to(device)
    q, r = torch.linalg.qr(random_matrix)
    q *= torch.sign(torch.diag(r)).unsqueeze(0)
    return q
//  [python]
def test_householder_product(self, device, dtype):
        def generate_reflectors_and_tau(A):
            """
            This function uses numpy.linalg.qr with mode "raw" to extract output of LAPACK's geqrf.
            There is torch.geqrf function but it doesn't work with complex-valued input.
            """
            if A.numel() > 0:
                A_cpu = A.cpu()
                flattened_batch_shape = [-1, *A_cpu.shape[-2:]]
                reflectors = torch.empty_like(A_cpu).view(*flattened_batch_shape)
                tau_shape = [*A_cpu.shape[:-2], A_cpu.shape[-1]]
                tau = torch.empty(tau_shape, dtype=dtype).view(-1, A_cpu.shape[-1])
                for A_i, reflectors_i, tau_i in zip(A_cpu.contiguous().view(*flattened_batch_shape), reflectors, tau):
                    reflectors_tmp, tau_i[:] = map(torch.from_numpy, np.linalg.qr(A_i, mode='raw'))
                    reflectors_i[:] = reflectors_tmp.T
                reflectors = reflectors.view(*A_cpu.shape)
                tau = tau.view(tau_shape)
                return reflectors.to(A.device), tau.to(A.device)

            reflectors = torch.empty_like(A)
            tau = torch.empty(*A.shape[:-2], A.shape[-1], dtype=dtype, device=device)
            return reflectors, tau

        def run_test(shape):
            A = torch.randn(*shape, dtype=dtype, device=device)
            reflectors, tau = generate_reflectors_and_tau(A)
            Q, R = torch.qr(A)
            expected = Q
            actual = torch.linalg.householder_product(reflectors, tau)
            # torch.linalg.qr does not work correctly for zero batch dimension tensors
            # see https://github.com/pytorch/pytorch/issues/50576
            if (A.numel() > 0):
                self.assertEqual(expected, actual)
            else:
                self.assertTrue(actual.shape == shape)

            # if tau is empty and A is not the result should be a matrix with ones on the diagonal
            if (A.numel() > 0):
                tau_empty = torch.empty(*shape[:-2], 0, dtype=dtype, device=device)
                identity_mat = torch.zeros_like(reflectors)
                identity_mat.diagonal(dim1=-1, dim2=-2)[:] = 1
                actual = torch.linalg.householder_product(reflectors, tau_empty)
                self.assertEqual(actual, identity_mat)

            out = torch.empty_like(A)
            ans = torch.linalg.householder_product(reflectors, tau, out=out)
            self.assertEqual(ans, out)
            if (A.numel() > 0):
                self.assertEqual(expected, out)

        shapes = [(0, 0), (5, 0),  # Empty matrix
                  (5, 5), (5, 3),  # Single matrix
                  (0, 0, 0), (0, 5, 5), (0, 5, 3),  # Zero batch dimension tensors
                  (2, 5, 5), (2, 5, 3),  # 3-dim tensors
                  (2, 1, 5, 5), (2, 1, 5, 3)]  # 4-dim tensors
        for shape in shapes:
            run_test(shape)
//  [python]
def test_householder_product(self, device, dtype):
        def generate_reflectors_and_tau(A):
            """
            This function uses numpy.linalg.qr with mode "raw" to extract output of LAPACK's geqrf.
            There is torch.geqrf function but it doesn't work with complex-valued input.
            """
            if A.numel() > 0:
                A_cpu = A.cpu()
                flattened_batch_shape = [-1, *A_cpu.shape[-2:]]
                reflectors = torch.empty_like(A_cpu).view(*flattened_batch_shape)
                tau_shape = [*A_cpu.shape[:-2], A_cpu.shape[-1]]
                tau = torch.empty(tau_shape, dtype=dtype).view(-1, A_cpu.shape[-1])
                for A_i, reflectors_i, tau_i in zip(A_cpu.contiguous().view(*flattened_batch_shape), reflectors, tau):
                    reflectors_tmp, tau_i[:] = map(torch.from_numpy, np.linalg.qr(A_i, mode='raw'))
                    reflectors_i[:] = reflectors_tmp.T
                reflectors = reflectors.view(*A_cpu.shape)
                tau = tau.view(tau_shape)
                return reflectors.to(A.device), tau.to(A.device)

            reflectors = torch.empty_like(A)
            tau = torch.empty(*A.shape[:-2], A.shape[-1], dtype=dtype, device=device)
            return reflectors, tau

        def run_test(shape):
            A = torch.randn(*shape, dtype=dtype, device=device)
            reflectors, tau = generate_reflectors_and_tau(A)
            expected, _ = torch.linalg.qr(A)
            actual = torch.linalg.householder_product(reflectors, tau)
            # torch.linalg.qr does not work correctly for zero batch dimension tensors
            # see https://github.com/pytorch/pytorch/issues/50576
            if (A.numel() > 0):
                self.assertEqual(expected, actual)
            else:
                self.assertTrue(actual.shape == shape)

            # if tau is empty and A is not the result should be a matrix with ones on the diagonal
            if (A.numel() > 0):
                tau_empty = torch.empty(*shape[:-2], 0, dtype=dtype, device=device)
                identity_mat = torch.zeros_like(reflectors)
                identity_mat.diagonal(dim1=-1, dim2=-2)[:] = 1
                actual = torch.linalg.householder_product(reflectors, tau_empty)
                self.assertEqual(actual, identity_mat)

            out = torch.empty_like(A)
            ans = torch.linalg.householder_product(reflectors, tau, out=out)
            self.assertEqual(ans, out)
            if (A.numel() > 0):
                self.assertEqual(expected, out)

        shapes = [(0, 0), (5, 0),  # Empty matrix
                  (5, 5), (5, 3),  # Single matrix
                  (0, 0, 0), (0, 5, 5), (0, 5, 3),  # Zero batch dimension tensors
                  (2, 5, 5), (2, 5, 3),  # 3-dim tensors
                  (2, 1, 5, 5), (2, 1, 5, 3)]  # 4-dim tensors
        for shape in shapes:
            run_test(shape)
//  [python]
def so_uniform_init_(tensor):
    r"""Fills in the input ``tensor`` in place with an orthogonal matrix.
    If square, the matrix will have positive determinant.
    The tensor will be distributed according to the Haar measure.
    The input tensor must have at least 2 dimensions.
    For tensors with more than 2 dimensions the first dimensions are treated as
    batch dimensions.
    Args:
        tensor (torch.Tensor): a 2-dimensional tensor or a batch of them
    """
    # We re-implement torch.nn.init.orthogonal_, as their treatment of batches
    # is not in a per-matrix base
    if tensor.ndim < 2:
        raise ValueError(
            "Only tensors with 2 or more dimensions are supported. "
            "Got a tensor of shape {}".format(tuple(tensor.size()))
        )
    n, k = tensor.size()[-2:]
    transpose = n < k
    with torch.no_grad():
        x = torch.empty_like(tensor).normal_(0, 1)
        if transpose:
            x.transpose_(-2, -1)
        q, r = torch.qr(x)

        # Make uniform (diag r >= 0)
        d = r.diagonal(dim1=-2, dim2=-1).sign()
        q *= d.unsqueeze(-2)
        if transpose:
            q.transpose_(-2, -1)

        # Make them have positive determinant by multiplying the
        # first column by -1 (does not change the measure)
        if n == k:
            mask = (torch.det(q) >= 0.0).float()
            mask[mask == 0.0] = -1.0
            mask = mask.unsqueeze(-1)
            q[..., 0] *= mask
        tensor.copy_(q)
        return tensor
//  [python]
def so_uniform_init_(tensor):
    r"""Fills in the input ``tensor`` in place with an orthogonal matrix.
    If square, the matrix will have positive determinant.
    The tensor will be distributed according to the Haar measure.
    The input tensor must have at least 2 dimensions.
    For tensors with more than 2 dimensions the first dimensions are treated as
    batch dimensions.
    Args:
        tensor (torch.Tensor): a 2-dimensional tensor or a batch of them
    """
    # We re-implement torch.nn.init.orthogonal_, as their treatment of batches
    # is not in a per-matrix base
    if tensor.ndim < 2:
        raise ValueError(
            "Only tensors with 2 or more dimensions are supported. "
            "Got a tensor of shape {}".format(tuple(tensor.size()))
        )
    n, k = tensor.size()[-2:]
    transpose = n < k
    with torch.no_grad():
        x = torch.empty_like(tensor).normal_(0, 1)
        if transpose:
            x.transpose_(-2, -1)
        q, r = torch.linalg.qr(x)

        # Make uniform (diag r >= 0)
        d = r.diagonal(dim1=-2, dim2=-1).sign()
        q *= d.unsqueeze(-2)
        if transpose:
            q.transpose_(-2, -1)

        # Make them have positive determinant by multiplying the
        # first column by -1 (does not change the measure)
        if n == k:
            mask = (torch.det(q) >= 0.0).float()
            mask[mask == 0.0] = -1.0
            mask = mask.unsqueeze(-1)
            q[..., 0] *= mask
        tensor.copy_(q)
        return tensor
//  [python]
def __init__(self, K=4, N=4, P=10, noise_power=1, episode_length=6, num_env=4, device=torch.device("cuda:0")):
        self.N = N  # #antennas
        self.K = K  # #users
        self.P = P  # Power
        self.noise_power = noise_power
        self.num_env = num_env
        self.device = device

        self.basis_vectors, _ = torch.qr(torch.rand(2 * self.K * self.N, 2 * self.K * self.N, dtype=torch.float))
        self.basis_vectors_tensor = self.basis_vectors.repeat(self.num_env, 1).reshape(self.num_env, 2 * self.K * self.N, 2 * self.K * self.N)
        self.subspace_dim = 2
        self.episode_length = episode_length
        self.generate_channel_tensor = functorch.vmap(self.generate_channel, in_dims=(0), randomness= 'different')
        self.step_tensor = functorch.vmap(self.step, in_dims=(0))
//  [python]
def __init__(self, K=4, N=4, P=10, noise_power=1, episode_length=6, num_env=4, device=torch.device("cuda:0")):
        self.N = N  # #antennas
        self.K = K  # #users
        self.P = P  # Power
        self.noise_power = noise_power
        self.num_env = num_env
        self.device = device

        self.basis_vectors, _ = torch.linalg.qr(torch.rand(2 * self.K * self.N, 2 * self.K * self.N, dtype=torch.float))
        self.basis_vectors_tensor = self.basis_vectors.repeat(self.num_env, 1).reshape(self.num_env, 2 * self.K * self.N, 2 * self.K * self.N)
        self.subspace_dim = 2
        self.episode_length = episode_length
        self.generate_channel_tensor = functorch.vmap(self.generate_channel, in_dims=(0), randomness= 'different')
        self.step_tensor = functorch.vmap(self.step, in_dims=(0))
//  [python]
def retraction(A, ref):
    data = A + ref
    Q, R = torch.qr(data)
    sign = (R.diag().sign() + 0.5).sign().diag()
    out = Q.mm(sign)
    return out
//  [python]
def retraction(A, ref):
    data = A + ref
    # Q, R = data.qr()
    Q,R = torch.linalg.qr(data)
    # To avoid (any possible) negative values in the output matrix, we multiply the negative values by -1
    sign = (R.diag().sign() + 0.5).sign().diag()
    out = Q.mm(sign)
    return out
//  [python]
def _synthesize_randomized_svd(
        self, k: int, p: int, q: int
    ) -> Tuple[Tensor, Tensor, Tensor]:
        r"""Synthesize eigendistortions using randomized truncated SVD.

        This method approximates the column space of the Fisher Info Matrix, projects the FIM into that column space,
        then computes its SVD.

        Parameters
        ----------
        k
            Number of eigenvecs (rank of factorization) to be returned.
        p
            Oversampling parameter, recommended to be 5.
        q
            Matrix power iteration. Used to squeeze the eigen spectrum for more accurate approximation.
            Recommended to be 2.

        Returns
        -------
        S
            Eigenvalues, Size((n, )).
        V
            Eigendistortions, Size((n, k)).
        error_approx
            Estimate of the approximation error. Defined as the expected error
            between the true subspace and approximated subspace.

        References
        -----
        [1] Halko, Martinsson, Tropp, Finding structure with randomness: Probabilistic algorithms for constructing
        approximate matrix decompositions, SIAM Rev. 53:2, pp. 217-288 https://arxiv.org/abs/0909.4061 (2011)

        """

        x, y = self._image_flat, self._representation_flat
        n = len(x)

        P = torch.randn(n, k + p).to(x.device)
        # orthogonalize first for numerical stability
        Q, R = torch.linalg.qr(P)
        P = Q
        _dummy_vec = torch.ones_like(y, requires_grad=True)
        Z = fisher_info_matrix_vector_product(y, x, P, _dummy_vec)

        # optional power iteration to squeeze the spectrum for more accurate
        # estimate
        for _ in range(q):
            Z = fisher_info_matrix_vector_product(y, x, Z, _dummy_vec)

        Q, _ = torch.linalg.qr(Z)
        # B = Q.T @ A @ Q
        B = Q.T @ fisher_info_matrix_vector_product(y, x, Q, _dummy_vec)
        _, S, Vh = torch.linalg.svd(B, False)  # eigendecomp of small matrix
        V = Vh.T
        V = Q @ V  # lift up to original dimensionality

        # estimate error in Q estimate of range space
        omega = fisher_info_matrix_vector_product(
            y, x, torch.randn(n, 20).to(x.device), _dummy_vec
        )
        error_approx = omega - (Q @ Q.T @ omega)
        error_approx = torch.linalg.vector_norm(error_approx, dim=0, ord=2).mean()

        return S[:k].clone(), V[:, :k].clone(), error_approx
//  [python]
def _synthesize_randomized_svd(
        self, k: int, p: int, q: int
    ) -> Tuple[Tensor, Tensor, Tensor]:
        r"""Synthesize eigendistortions using randomized truncated SVD.

        This method approximates the column space of the Fisher Info Matrix, projects the FIM into that column space,
        then computes its SVD.

        Parameters
        ----------
        k
            Number of eigenvecs (rank of factorization) to be returned.
        p
            Oversampling parameter, recommended to be 5.
        q
            Matrix power iteration. Used to squeeze the eigen spectrum for more accurate approximation.
            Recommended to be 2.

        Returns
        -------
        S
            Eigenvalues, Size((n, )).
        V
            Eigendistortions, Size((n, k)).
        error_approx
            Estimate of the approximation error. Defined as the expected error
            between the true subspace and approximated subspace.

        References
        -----
        [1] Halko, Martinsson, Tropp, Finding structure with randomness: Probabilistic algorithms for constructing
        approximate matrix decompositions, SIAM Rev. 53:2, pp. 217-288 https://arxiv.org/abs/0909.4061 (2011)

        """

        x, y = self._image_flat, self._representation_flat
        n = len(x)

        P = torch.randn(n, k + p).to(x.device)
        # orthogonalize first for numerical stability
        P, _ = torch.linalg.qr(P, "reduced")
        _dummy_vec = torch.ones_like(y, requires_grad=True)
        Z = fisher_info_matrix_vector_product(y, x, P, _dummy_vec)

        # optional power iteration to squeeze the spectrum for more accurate
        # estimate
        for _ in range(q):
            Z = fisher_info_matrix_vector_product(y, x, Z, _dummy_vec)

        Q, _ = torch.linalg.qr(Z, "reduced")
        # B = Q.T @ A @ Q
        B = Q.T @ fisher_info_matrix_vector_product(y, x, Q, _dummy_vec)
        _, S, Vh = torch.linalg.svd(B, False)  # eigendecomp of small matrix
        V = Vh.T
        V = Q @ V  # lift up to original dimensionality

        # estimate error in Q estimate of range space
        omega = fisher_info_matrix_vector_product(
            y, x, torch.randn(n, 20).to(x.device), _dummy_vec
        )
        error_approx = omega - (Q @ Q.T @ omega)
        error_approx = torch.linalg.vector_norm(error_approx, dim=0, ord=2).mean()

        return S[:k].clone(), V[:, :k].clone(), error_approx
//  [python]
def basis(A):
    """Return orthogonal basis of A columns."""
    Q, R = torch.qr(A)
    return Q
//  [python]
def basis(A):
    """Return orthogonal basis of A columns."""
    return torch.linalg.qr(A).Q
//  [python]
def _build_data_matricies(self):
        """
        The data matrix is build column-wise to represent regression matrix for least
        squares solver. The data matrix is QR decomposed.
        """

        self._b = self._data[self._order:]

        cols = [torch.ones(self._number_of_observations - self._order)]
        for i in range(self._order):
            col = self._data[(self._order - 1) - i: self._number_of_observations - i - 1]
            cols.append(col)

        A = torch.column_stack(cols)
        self._Q, self._R = torch.qr(A)
//  [python]
def _build_data_matricies(self):
        """
        The data matrix is build column-wise to represent regression matrix for least
        squares solver. The data matrix is QR decomposed.
        """

        self._b = self._data[self._order:]

        cols = [torch.ones(self._number_of_observations - self._order)]
        for i in range(self._order):
            col = self._data[(self._order - 1) - i: self._number_of_observations - i - 1]
            cols.append(col)

        A = torch.column_stack(cols)
        self._Q, self._R = torch.linalg.qr(A)
//  [python]
def orthogonal(rows, cols, gain=1):
    r"""return 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).
    Args:
        rows: rows
        cols: cols
        gain: optional scaling factor
    Examples:
        >>> orthogonal_(5, 3)
    """
    flattened = torch.randn((rows, cols))
    if rows < cols: flattened = 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)
    q *= d.sign()

    if rows < cols: q = q.T
    q *= gain
    return q
//  [python]
def orthogonal(rows, cols, gain=1):
    r"""return 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).
    Args:
        rows: rows
        cols: cols
        gain: optional scaling factor
    Examples:
        >>> orthogonal_(5, 3)
    """
    flattened = torch.randn((rows, cols))
    if rows < cols: flattened = 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)
    q *= d.sign()

    if rows < cols: q = q.T
    q *= gain
    return q
//  [python]
def local_encoder_pullback_xt(
            self, x=None, t=None, op=None, block_idx=None, pca_rank=16, chunk_size=25,
            min_iter=10, max_iter=100, convergence_threshold=1e-3,
        ):
        '''
        Args
            - sample : zt
            - op : ['down', 'mid', 'up']
            - block_idx : op == down, up : [0,1,2,3], op == mid : [0]
            - pooling : ['pixel-sum', 'channel-sum', 'single-channel', 'multiple-channel']
        Returns
            - h : hidden feature
        '''
        # necessary variables
        num_chunk = pca_rank // chunk_size if pca_rank % chunk_size == 0 else pca_rank // chunk_size + 1
        get_h = lambda x : self.get_h(
            x, t=t, op=op, block_idx=block_idx,
        )
        h_shape = get_h(x).shape
        print('h_shape : ', h_shape)

        c_i, w_i, h_i = x.size(1), x.size(2), x.size(3)
        c_o, w_o, h_o = h_shape[1], h_shape[2], h_shape[3]

        a = th.tensor(0., device=x.device)

        # Algorithm 1
        vT = th.randn(c_i*w_i*h_i, pca_rank, device=x.device)
        q, r = th.qr(vT)
        vT = q
        v = vT.T
        v = v.view(-1, c_i, w_i, h_i)

        for i in range(max_iter):
            v_prev = v.detach().cpu().clone()

            u = []
            time_s = time.time()

            v_buffer = list(v.chunk(num_chunk))
            for vi in v_buffer:
                # g = lambda a : get_h(x + a*vi.unsqueeze(0) if vi.size(0) == v.size(-1) else x + a*vi)
                g = lambda a : get_h(x + a*vi)
                ui = th.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
                u.append(ui.detach().cpu().clone())
            time_e = time.time()
            print('single v jacfwd t ==', time_e - time_s)
            u = th.cat(u, dim=0)
            u = u.to(x.device)

            # time_s = time.time()
            # g = lambda a : get_h(x + a*v)
            # u = th.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
            # time_e = time.time()
            # print('single vi jacfwd t ==', time_e - time_s)

            g = lambda x : einsum(u, get_h(x), 'b c w h, i c w h -> b')
            v_ = th.autograd.functional.jacobian(g, x)
            v_ = v_.view(-1, c_i*w_i*h_i)

            _, s, v = th.linalg.svd(v_, full_matrices=False)
            v = v.view(-1, c_i, w_i, h_i)
            u = u.view(-1, c_o, w_o, h_o)
            
            convergence = th.dist(v_prev, v.detach().cpu())
            print(f'power method : {i}-th step convergence : ', convergence)
            if th.allclose(v_prev, v.detach().cpu(), atol=convergence_threshold) and (i > min_iter):
                print('reach convergence threshold : ', convergence)
                break

            if i == max_iter - 1:
                print('last convergence : ', convergence)

        u, s, vT = u.view(-1, c_o*w_o*h_o).T.detach(), s.sqrt().detach(), v.view(-1, c_i*w_i*h_i).detach()
        return u, s, vT
//  [python]
def local_encoder_pullback_xt(
            self, x=None, t=None, op=None, block_idx=None, pca_rank=16, chunk_size=25,
            min_iter=10, max_iter=100, convergence_threshold=1e-3,
        ):
        '''
        Args
            - sample : zt
            - op : ['down', 'mid', 'up']
            - block_idx : op == down, up : [0,1,2,3], op == mid : [0]
            - pooling : ['pixel-sum', 'channel-sum', 'single-channel', 'multiple-channel']
        Returns
            - h : hidden feature
        '''
        # necessary variables
        num_chunk = pca_rank // chunk_size if pca_rank % chunk_size == 0 else pca_rank // chunk_size + 1
        get_h = lambda x : self.get_h(
            x, t=t, op=op, block_idx=block_idx,
        )
        h_shape = get_h(x).shape
        print('h_shape : ', h_shape)

        c_i, w_i, h_i = x.size(1), x.size(2), x.size(3)
        c_o, w_o, h_o = h_shape[1], h_shape[2], h_shape[3]

        a = th.tensor(0., device=x.device)

        # Algorithm 1
        vT = th.randn(c_i*w_i*h_i, pca_rank, device=x.device)
        vT, _ = th.linalg.qr(vT)
        v = vT.T
        v = v.view(-1, c_i, w_i, h_i)

        for i in range(max_iter):
            v_prev = v.detach().cpu().clone()

            u = []
            time_s = time.time()

            v_buffer = list(v.chunk(num_chunk))
            for vi in v_buffer:
                # g = lambda a : get_h(x + a*vi.unsqueeze(0) if vi.size(0) == v.size(-1) else x + a*vi)
                g = lambda a : get_h(x + a*vi)
                ui = th.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
                u.append(ui.detach().cpu().clone())
            time_e = time.time()
            print('single v jacfwd t ==', time_e - time_s)
            u = th.cat(u, dim=0)
            u = u.to(x.device)

            # time_s = time.time()
            # g = lambda a : get_h(x + a*v)
            # u = th.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
            # time_e = time.time()
            # print('single vi jacfwd t ==', time_e - time_s)

            g = lambda x : einsum(u, get_h(x), 'b c w h, i c w h -> b')
            v_ = th.autograd.functional.jacobian(g, x)
            v_ = v_.view(-1, c_i*w_i*h_i)

            _, s, v = th.linalg.svd(v_, full_matrices=False)
            v = v.view(-1, c_i, w_i, h_i)
            u = u.view(-1, c_o, w_o, h_o)
            
            convergence = th.dist(v_prev, v.detach().cpu())
            print(f'power method : {i}-th step convergence : ', convergence)
            if th.allclose(v_prev, v.detach().cpu(), atol=convergence_threshold) and (i > min_iter):
                print('reach convergence threshold : ', convergence)
                break

            if i == max_iter - 1:
                print('last convergence : ', convergence)

        u, s, vT = u.view(-1, c_o*w_o*h_o).T.detach(), s.sqrt().detach(), v.view(-1, c_i*w_i*h_i).detach()
        return u, s, vT
//  [python]
def lr(x1, x2):
    x1_flat, x2_flat = rearrange_activations(x1), rearrange_activations(x2)
    q2, r2 = torch.qr(x2_flat)
    return (torch.linalg.norm(q2.T @ x1_flat)) ** 2 / (torch.linalg.norm(x1_flat)) ** 2
//  [python]
def lr(x1, x2):
    x1_flat, x2_flat = rearrange_activations(x1), rearrange_activations(x2)
    q2, r2 = torch.linalg.qr(x2_flat)
    return (torch.linalg.norm(q2.T @ x1_flat)) ** 2 / (torch.linalg.norm(x1_flat)) ** 2
//  [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 = torch.linalg.qr(unstructured_block.cpu(), 'reduced')
    q, r = map(lambda t: t.to(device), (q, r))
    return q.t()
//  [python]
def update_kernel_normal_bases(self):
        if self.Policy.n_kernels > 0:
            # reset tensors for current bases
            self.kernel_obstacle_bases_tmp *= 0
            self.kernel_obstacle_bases_tmp[0:self.Policy.n_kernels][:] += torch.eye(self.n_dof)

            # calculate gamma and repulsion for kernel centers
            dst, grad = self.distance_repulsion_nn(self.Policy.mu_c[0:self.Policy.n_kernels], aot=False)
            self.kernel_obstacle_bases_tmp[0:self.Policy.n_kernels][:, :, 0] = grad

            #perform qr decomposition
            E, R = torch.qr(self.kernel_obstacle_bases_tmp[0:self.Policy.n_kernels].cpu().to(torch.float32))
            E = E.to(**self.tensor_args)
            E[:, :, 0] = grad / grad.norm(2, 1).unsqueeze(1)
            # print((self.Policy.kernel_obstacle_bases[0:self.Policy.n_kernels] - E).norm(p=2)) #debug update difference

            # update policies kernel bases
            self.Policy.kernel_obstacle_bases[0:self.Policy.n_kernels] = E
        else:
            pass
        return 0
//  [python]
def update_kernel_normal_bases(self):
        if self.Policy.n_kernels > 0:
            # reset tensors for current bases
            self.kernel_obstacle_bases_tmp *= 0
            self.kernel_obstacle_bases_tmp[0:self.Policy.n_kernels][:] += torch.eye(self.n_dof)

            # calculate gamma and repulsion for kernel centers
            dst, grad = self.distance_repulsion_nn(self.Policy.mu_c[0:self.Policy.n_kernels], aot=False)
            self.kernel_obstacle_bases_tmp[0:self.Policy.n_kernels][:, :, 0] = grad

            #perform qr decomposition
            E, R = torch.linalg.qr(self.kernel_obstacle_bases_tmp[0:self.Policy.n_kernels].cpu().to(torch.float32))
            E = E.to(**self.tensor_args)
            E[:, :, 0] = grad / grad.norm(2, 1).unsqueeze(1)
            # print((self.Policy.kernel_obstacle_bases[0:self.Policy.n_kernels] - E).norm(p=2)) #debug update difference

            # update policies kernel bases
            self.Policy.kernel_obstacle_bases[0:self.Policy.n_kernels] = E
        else:
            pass
        return 0
//  [python]
def create_random_ortho_matrix(n, seed=-1):
    dtype = torch.float64
    if seed > 0:
        torch.manual_seed(seed)
    a = torch.randn((n, n), dtype=dtype)
    q, r = torch.qr(a)
    return q
//  [python]
def create_random_ortho_matrix(n, seed=-1):
    dtype = torch.float64
    if seed > 0:
        torch.manual_seed(seed)
    a = torch.randn((n, n), dtype=dtype)
    q, r = torch.linalg.qr(a)
    return q
//  [python]
def QR(mat):
    """
    Compute the QR decomposition. Backend can be changed.

    Parameters
    ----------
    mat : tn array
        DESCRIPTION.

    Returns
    -------
    Q : the Q matrix
    R : the R matrix

    """
    Q,R = t.linalg.qr(mat)
    return Q, R
//  [python]
def QR(mat):
    """
    Compute the QR decomposition. Backend can be changed.

    Parameters
    ----------
    mat : tn array
        DESCRIPTION.

    Returns
    -------
    Q : the Q matrix
    R : the R matrix

    """
    Q,R = tn.linalg.qr(mat)
    return Q, R
//  [python]
def randomised_SVD_jacobian(input, model, n_rank):

    model.eval()
    forward_maps = []
    for _ in range(n_rank):
        forward_maps.append(central_diff_approx(input, model))
    forward_map = torch.cat(forward_maps).transpose(-2, -1)
    q, _ = torch.qr(forward_map)
    b_t_ = []
    for l in range(q.size()[-1]):
        q_l = q[:, l].detach()
        out = model(input) @ q_l
        model.zero_grad()
        out.backward()
        b_l = agregate_flatten_weight_grad(model).detach()
        b_t_.append(b_l)
    b_matrix = torch.stack(b_t_, dim=0)
    u, s, vh = torch.svd_lowrank(b_matrix, q=n_rank, niter=2, M=None)
    v = vh.transpose(-2, -1).conj()
    return s, v
//  [python]
def randomised_SVD_jacobian(input, model, n_rank):

    model.eval()
    forward_maps = []
    for _ in range(n_rank):
        forward_maps.append(central_diff_approx(input, model))
    forward_map = torch.cat(forward_maps).transpose(-2, -1)
    q, _ = torch.linalg.qr(forward_map)
    b_t_ = []
    for l in range(q.size()[-1]):
        q_l = q[:, l].detach()
        out = model(input) @ q_l
        model.zero_grad()
        out.backward()
        b_l = agregate_flatten_weight_grad(model).detach()
        b_t_.append(b_l)
    b_matrix = torch.stack(b_t_, dim=0)
    u, s, vh = torch.svd_lowrank(b_matrix, q=n_rank, niter=2, M=None)
    v = vh.transpose(-2, -1).conj()
    return s, v
//  [python]
def get_new_K(K1, K2, K3):
    q1, r1 = torch.qr(K1)
    q3, r3 = torch.qr(torch.transpose(K3, -1, -2))
    middle_k = full_tt(r1, K2, torch.transpose(r3, -1, -2))
    return q1, middle_k, torch.transpose(q3, -1, -2)
//  [python]
def get_new_K(K1, K2, K3):
    # (m, R) -> (m, min(m, R)) -- q: q.T @ q = I, (min(m, R), R) -- r
    q1, r1 = torch.linalg.qr(K1)
    q3, r3 = torch.linalg.qr(torch.transpose(K3, -1, -2))
    middle_k = full_tt(r1, K2, torch.transpose(r3, -1, -2))
    return q1, middle_k, torch.transpose(q3, -1, -2)
//  [python]
def __init__(self, channels, n_split=4, no_jacobian=False, lu=True, n_sqz=2, **kwargs):
        super().__init__()
        assert (n_split % 2 == 0)
        self.channels = channels
        self.n_split = n_split
        self.n_sqz = n_sqz
        self.no_jacobian = no_jacobian

        w_init = torch.linalg.qr(torch.FloatTensor(self.n_split, self.n_split).normal_(), 'complete')[0]
        if torch.det(w_init) < 0:
            w_init[:, 0] = -1 * w_init[:, 0]
        self.lu = lu
        if lu:
            # LU decomposition can slightly speed up the inverse
            np_p, np_l, np_u = scipy.linalg.lu(w_init)
            np_s = np.diag(np_u)
            np_sign_s = np.sign(np_s)
            np_log_s = np.log(np.abs(np_s))
            np_u = np.triu(np_u, k=1)
            l_mask = np.tril(np.ones(w_init.shape, dtype=float), -1)
            eye = np.eye(*w_init.shape, dtype=float)

            self.register_buffer('p', torch.Tensor(np_p.astype(float)))
            self.register_buffer('sign_s', torch.Tensor(np_sign_s.astype(float)))
            self.l = nn.Parameter(torch.Tensor(np_l.astype(float)), requires_grad=True)
            self.log_s = nn.Parameter(torch.Tensor(np_log_s.astype(float)), requires_grad=True)
            self.u = nn.Parameter(torch.Tensor(np_u.astype(float)), requires_grad=True)
            self.register_buffer('l_mask', torch.Tensor(l_mask))
            self.register_buffer('eye', torch.Tensor(eye))
        else:
            self.weight = nn.Parameter(w_init)