jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def local_encoder_pullback_xt(
        self, x, t, op=None, block_idx=None,
        pca_rank=50, chunk_size=25, min_iter=10, max_iter=100, convergence_threshold=1e-3,
    ):
    '''
    Args
        - x : 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
    '''
    num_chunk = pca_rank // chunk_size if pca_rank % chunk_size == 0 else pca_rank // chunk_size + 1

    # get h samples
    time_s = time.time()

    # necessary variables
    h_shape = self.get_h(
        x, t=t, op=op, block_idx=block_idx, 
    ).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 = torch.tensor(0., device=x.device, dtype=x.dtype)

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

    time_s = time.time()
    for i in range(max_iter):
        v = v.to(device=x.device, dtype=x.dtype)
        v_prev = v.detach().cpu().clone()
        
        
        u = []
        v_buffer = list(v.chunk(num_chunk))
        for vi in v_buffer:
            # g = lambda a : get_h(x + a*vi)

            g = lambda a : self.get_h(
                x + a*vi, t=t, op=op, block_idx=block_idx, 
            )

            ui = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
            u.append(ui.detach().cpu().clone())
        u = torch.cat(u, dim=0)
        u = u.to(x.device, x.dtype)

        # time_s = time.time()
        # g = lambda a : get_h(x + a*v)
        # u = torch.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, self.get_h(
                x, t=t, op=op, block_idx=block_idx, 
            ), 'b c w h, i c w h -> b'
        )
        v_ = torch.autograd.functional.jacobian(g, x)
        v_ = v_.view(-1, c_i*w_i*h_i)

        _, s, v = torch.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 = torch.dist(v_prev, v.detach().cpu()).item()
        print(f'power method : {i}-th step convergence : ', convergence)
        
        if torch.allclose(v_prev, v.detach().cpu(), atol=convergence_threshold) and (i > min_iter):
            print('reach convergence threshold : ', convergence)
            break

    time_e = time.time()
    print('power method runtime ==', time_e - time_s)

    # 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()
    u, s, vT = u.reshape(-1, c_o*w_o*h_o).T.detach(), s.sqrt().detach(), v.reshape(-1, c_i*w_i*h_i).detach()
    return u, s, vT
//  [python]
def fake_poweriteration_with_outlierquant(
    input: torch.Tensor, loop, rank, device, p_base, q_base, quantize_bit, left
):
    batch, num_head, seq_len, sep_dim = input.shape
    input = (
        input.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )
    input = input.float()
    # convert to 32bits for qr decomposition
    input = input.view(seq_len * batch, sep_dim * num_head)
    # 3 calculation = loop * (matmul) + 2 * qrO(n^2)
    for i in range(loop):
        if i == loop - 1:
            q, r = torch.qr(p_base[0])
            p_base[0] = q
        q_base[0] = input @ p_base[0]
        if i == loop - 1:
            q, r = torch.qr(q_base[0])
            q_base[0] = q
        p_base[0] = torch.transpose(input, 0, 1) @ q_base[0]
    # outlier quantization
    q_base[0] = fake_dense_sparse_uniformquantization(q_base[0], quantize_bit, left)
    p_base[0] = fake_dense_sparse_uniformquantization(p_base[0], quantize_bit, left)
    input = q_base[0] @ torch.transpose(p_base[0], 0, 1)
    input = input.view(batch, seq_len, num_head, sep_dim)
    input = input.permute(0, 2, 1, 3)
    input = input.half()
    return input
//  [python]
def fake_poweriteration_with_outlierquant(
    input: torch.Tensor, loop, rank, device, p_base, q_base, quantize_bit, left
):
    batch, num_head, seq_len, sep_dim = input.shape
    input = (
        input.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )
    input = input.float()
    # convert to 32bits for qr decomposition
    input = input.view(seq_len * batch, sep_dim * num_head)
    # 3 calculation = loop * (matmul) + 2 * qrO(n^2)
    for i in range(loop):
        if i == loop - 1:
            p_base[0] = torch.linalg.qr(p_base[0]).Q
        q_base[0] = input @ p_base[0]
        if i == loop - 1:
            q_base[0] = torch.linalg.qr(q_base[0]).Q
        p_base[0] = torch.transpose(input, 0, 1) @ q_base[0]
    # outlier quantization
    q_base[0] = fake_dense_sparse_uniformquantization(q_base[0], quantize_bit, left)
    p_base[0] = fake_dense_sparse_uniformquantization(p_base[0], quantize_bit, left)
    input = q_base[0] @ torch.transpose(p_base[0], 0, 1)
    input = input.view(batch, seq_len, num_head, sep_dim)
    input = input.permute(0, 2, 1, 3)
    input = input.half()
    return input
//  [python]
def local_decoder_pullback_zt(
        self, sample, timestep, encoder_hidden_states=None, op=None, block_idx=None,
        pca_rank=50, chunk_size=25, min_iter=10, max_iter=100, convergence_threshold=None,
    ):
    '''
    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
    '''
    # get h samples
    time_s = time.time()

    # necessary variables
    h = self.get_h(
        sample, timestep=timestep, encoder_hidden_states=encoder_hidden_states,
        op=op, block_idx=block_idx, 
    )

    get_h_to_e = lambda h : self.get_h_to_e(
        sample=sample, timestep=timestep, encoder_hidden_states=encoder_hidden_states,
        input_h=h, op=op, block_idx=block_idx, verbose=False,
    )

    c_i, w_i, h_i = h.size(1), h.size(2), h.size(3)
    c_o, w_o, h_o = sample.size(1), sample.size(2), sample.size(3)

    a = torch.tensor(0., device=sample.device, dtype=sample.dtype)

    # Algorithm 1
    vT = torch.randn(c_i*w_i*h_i, pca_rank, device=sample.device, dtype=torch.float)
    q, r = torch.linalg.qr(vT)
    v = q.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 = v.to(device=sample.device, dtype=sample.dtype)
        v_buffer = list(v.chunk(v.size(0) // chunk_size))
        for vi in v_buffer:
            g = lambda a : get_h_to_e(h + a*vi)
            ui = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
            u.append(ui.detach().cpu().clone())
        u = torch.cat(u, dim=0)
        u = u.to(sample.device, sample.dtype)
        time_e = time.time()
        print('single v jacfwd t ==', time_e - time_s)

        # time_s = time.time()
        # g = lambda a : get_h(x + a*v)
        # u = torch.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 h : einsum(u, get_h_to_e(h), 'b c w h, i c w h -> b')
        v_ = torch.autograd.functional.jacobian(g, h)
        v_ = v_.view(-1, c_i*w_i*h_i)

        _, s, v = torch.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 = torch.dist(v_prev, v.detach().cpu()).item()
        print(f'power method : {i}-th step convergence : ', convergence)

        if torch.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)

    # decoder jac do not return vT
    u, s, vT = v.view(-1, c_i*w_i*h_i).T.detach(), s.sqrt().detach(), u.view(-1, c_o*w_o*h_o).detach()
    return u, s, vT
//  [python]
def local_decoder_pullback_zt(
        self, sample, timestep, encoder_hidden_states=None, op=None, block_idx=None,
        pca_rank=50, chunk_size=25, min_iter=10, max_iter=100, convergence_threshold=None,
    ):
    '''
    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
    '''
    # get h samples
    time_s = time.time()

    # necessary variables
    h = self.get_h(
        sample, timestep=timestep, encoder_hidden_states=encoder_hidden_states,
        op=op, block_idx=block_idx, 
    )

    get_h_to_e = lambda h : self.get_h_to_e(
        sample=sample, timestep=timestep, encoder_hidden_states=encoder_hidden_states,
        input_h=h, op=op, block_idx=block_idx, verbose=False,
    )

    c_i, w_i, h_i = h.size(1), h.size(2), h.size(3)
    c_o, w_o, h_o = sample.size(1), sample.size(2), sample.size(3)

    a = torch.tensor(0., device=sample.device, dtype=sample.dtype)

    # Algorithm 1
    vT = torch.randn(c_i*w_i*h_i, pca_rank, device=sample.device, dtype=torch.float)
    vT, _ = torch.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 = v.to(device=sample.device, dtype=sample.dtype)
        v_buffer = list(v.chunk(v.size(0) // chunk_size))
        for vi in v_buffer:
            g = lambda a : get_h_to_e(h + a*vi)
            ui = torch.func.jacfwd(g, argnums=0, has_aux=False, randomness='error')(a)
            u.append(ui.detach().cpu().clone())
        u = torch.cat(u, dim=0)
        u = u.to(sample.device, sample.dtype)
        time_e = time.time()
        print('single v jacfwd t ==', time_e - time_s)

        # time_s = time.time()
        # g = lambda a : get_h(x + a*v)
        # u = torch.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 h : einsum(u, get_h_to_e(h), 'b c w h, i c w h -> b')
        v_ = torch.autograd.functional.jacobian(g, h)
        v_ = v_.view(-1, c_i*w_i*h_i)

        _, s, v = torch.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 = torch.dist(v_prev, v.detach().cpu()).item()
        print(f'power method : {i}-th step convergence : ', convergence)

        if torch.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)

    # decoder jac do not return vT
    u, s, vT = v.view(-1, c_i*w_i*h_i).T.detach(), s.sqrt().detach(), u.view(-1, c_o*w_o*h_o).detach()
    return u, s, vT
//  [python]
def ptcompress(cache: torch.Tensor, p_buffer, q_buffer, loop):
    # cache size [batch,num_head,seq_len,model_dim/num_head]
    # -> [batch,seq_len,model_dim] -> [batch * seq_len,model_dim]
    # p_buffer of size [model_dim, rank], q_buffer of size [batch * seq_len,rank]
    batch, num_head, seq_len, sep_dim = cache.shape
    cache = (
        cache.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )  # convert to 32bits for qr decomposition
    p_buffer[0] = p_buffer[0]
    q_buffer[0] = q_buffer[0]
    cache = cache.view(seq_len * batch, sep_dim * num_head)
    for i in range(loop):
        if i == loop - 1:
            p_buffer[0] = torch.qr(p_buffer[0]).Q
        q_buffer[0] = cache @ p_buffer[0]
        if i == loop - 1:
            q_buffer[0] = torch.qr(q_buffer[0]).Q
        p_buffer[0] = torch.transpose(cache, 0, 1) @ q_buffer[0]
    p_buffer[0] = p_buffer[0]
    q_buffer[0] = q_buffer[0]
    cache = cache.view(batch, num_head, seq_len, sep_dim)
    return p_buffer, q_buffer, batch, num_head, seq_len, sep_dim
//  [python]
def ptcompress(cache: torch.Tensor, p_buffer, q_buffer, loop):
    # cache size [batch,num_head,seq_len,model_dim/num_head]
    # -> [batch,seq_len,model_dim] -> [batch * seq_len,model_dim]
    # p_buffer of size [model_dim, rank], q_buffer of size [batch * seq_len,rank]
    batch, num_head, seq_len, sep_dim = cache.shape
    cache = (
        cache.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )  # convert to 32bits for qr decomposition
    p_buffer[0] = p_buffer[0]
    q_buffer[0] = q_buffer[0]
    cache = cache.view(seq_len * batch, sep_dim * num_head)
    for i in range(loop):
        if i == loop - 1:
            p_buffer[0] = torch.linalg.qr(p_buffer[0]).Q
        q_buffer[0] = cache @ p_buffer[0]
        if i == loop - 1:
            q_buffer[0] = torch.linalg.qr(q_buffer[0]).Q
        p_buffer[0] = torch.transpose(cache, 0, 1) @ q_buffer[0]
    p_buffer[0] = p_buffer[0]
    q_buffer[0] = q_buffer[0]
    cache = cache.view(batch, num_head, seq_len, sep_dim)
    return p_buffer, q_buffer, batch, num_head, seq_len, sep_dim
//  [python]
def fake_poweriteration(input: torch.Tensor, loop, rank, device, p_base, q_base):
    # input size [batch,num_head,seq_len,model_dim/num_head]
    # -> [batch,seq_len,model_dim] -> [batch * seq_len,model_dim]
    # p_base = torch.rand(input.shape[3] * input.shape[1], rank).to(device)
    # q_base = torch.rand(input.shape[0] * input.shape[2], rank).to(device)
    batch, num_head, seq_len, sep_dim = input.shape
    input = (
        input.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )  # convert to 32bits for qr decomposition
    input = input.view(seq_len * batch, sep_dim * num_head)
    input = input.float()
    # print(p_base)
    p_base[0] = p_base[0].float()
    q_base[0] = q_base[0].float()
    # 3 calculation = loop * (matmul) + 2 * qrO(n^2)
    for i in range(loop):
        if i == loop - 1:
            q, r = torch.qr(p_base[0])
            p_base[0] = q
        q_base[0] = input @ p_base[0]
        if i == loop - 1:
            q, r = torch.qr(q_base[0])
            q_base[0] = q
        p_base[0] = torch.transpose(input, 0, 1) @ q_base[0]
    input = q_base[0] @ torch.transpose(p_base[0], 0, 1)
    input = input.view(batch, seq_len, num_head, sep_dim)
    input = input.permute(0, 2, 1, 3)
    input = input.half()
    p_base[0] = p_base[0].half()
    q_base[0] = q_base[0].half()
    return input
//  [python]
def fake_poweriteration(input: torch.Tensor, loop, rank, device, p_base, q_base):
    # input size [batch,num_head,seq_len,model_dim/num_head]
    # -> [batch,seq_len,model_dim] -> [batch * seq_len,model_dim]
    # p_base = torch.rand(input.shape[3] * input.shape[1], rank).to(device)
    # q_base = torch.rand(input.shape[0] * input.shape[2], rank).to(device)
    batch, num_head, seq_len, sep_dim = input.shape
    input = (
        input.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )  # convert to 32bits for qr decomposition
    input = input.view(seq_len * batch, sep_dim * num_head)
    input = input.float()
    # print(p_base)
    p_base[0] = p_base[0].float()
    q_base[0] = q_base[0].float()
    # 3 calculation = loop * (matmul) + 2 * qrO(n^2)
    for i in range(loop):
        if i == loop - 1:
            p_base[0] = torch.linalg.qr(p_base[0]).Q
        q_base[0] = input @ p_base[0]
        if i == loop - 1:
            q_base[0] = torch.linalg.qr(q_base[0]).Q
        p_base[0] = torch.transpose(input, 0, 1) @ q_base[0]
    input = q_base[0] @ torch.transpose(p_base[0], 0, 1)
    input = input.view(batch, seq_len, num_head, sep_dim)
    input = input.permute(0, 2, 1, 3)
    input = input.half()
    p_base[0] = p_base[0].half()
    q_base[0] = q_base[0].half()
    return input
//  [python]
def find_d(Z, z_c, r_dimcheck, n_iter_inner, d_prev, max_error, max_error_dimcheck_ratio=0.3):
    # We find the minimial d by iteratively fitting a model
    # for some size d, then increase d and repeat if the max
    # reconstruction error is too large

    max_error = max_error_dimcheck_ratio * r_dimcheck

    N, D = Z.shape
    # init
    U_0 = torch.randn(D, d_prev)
    q, r = torch.linalg.qr(U_0)
    U_0 = q

    # TRACKING VARS
    was_decreasing = False

    # init tracking variable
    # max_error = max_error_ratio*r_dimcheck

    # note kernel will stay the same for all d
    gamma = float(np.log(2)) / (r_dimcheck ** 2)
    kernel_pre = torch.exp(-gamma * (Z - z_c).pow(2).sum(dim=1, keepdim=True))
    z_mu = (Z * kernel_pre).sum(dim=0, keepdim=True) / kernel_pre.sum()
    kernel = torch.exp(-gamma * (Z - z_mu).pow(2).sum(dim=1, keepdim=True))

    # need to track U_prev in case of decreasing
    U_prev = U_0.clone()

    # note this loop iterates a maximum of D times
    for i in range(D):
        U, V, loss = opt_UV(Z, z_c, U_0, n_iter_inner, kernel=kernel)

        # compute max error, if too large, increase d
        coord = (Z - z_c) @ U
        Z_perp = (Z - z_c) - coord @ U.T
        # NOTE THIS NEEDS TO CHANGE IF DOING DIAG HESSIAN
        H_input = torch.bmm(coord.unsqueeze(2), coord.unsqueeze(1))
        # since tensor is symmetric, we only need to keep upper diag terms
        # output is a tensor of shape (N, d(d+1)/2)
        d = U.shape[1]
        idx_triu = torch.triu_indices(d, d)
        H_input = H_input[:, idx_triu[0, :], idx_triu[1, :]]
        l0_error = (kernel * (Z_perp - H_input @ V.T)).norm(dim=1).max()
        # l2_error = 0.5*(kernel * (Z_perp - H_input @ V.T)).norm(dim=1).pow(2).mean()
        # if achieved desired loss, reduce dimension
        # print(f'd: {d}, error: {l0_error}')
        if l0_error <= max_error:
            # if we were increasing but then achieved desired loss we have converged
            if (not was_decreasing and i > 0) or U.shape[1] == 1:
                break

            # otherwise, continue to reduce dimension
            was_decreasing = True
            # remove column of U lease correlated with V
            idx_min = ((Z - z_c) @ U).abs().mean(dim=0).argmin()
            U_0 = torch.cat((U[:, :idx_min], U[:, idx_min + 1:]), dim=1)
            U_prev = U.clone()

        else:
            # if we were decreasing but then achieved desired loss we have converged
            if was_decreasing or U.shape[1] == D:
                U = U_prev
                break
            U_new = torch.randn(D, 1)
            U_new = U_new - U @ (U.T @ U_new)
            U_new = U_new / torch.norm(U_new, p=2)
            U_0 = torch.cat((U, U_new), dim=1).clone()

    return U.detach().data, loss.detach().data
//  [python]
def find_d(Z, z_c, r_dimcheck, n_iter_inner, d_prev, max_error, max_error_dimcheck_ratio=0.3):
    # We find the minimial d by iteratively fitting a model
    # for some size d, then increase d and repeat if the max
    # reconstruction error is too large

    max_error = max_error_dimcheck_ratio * r_dimcheck

    N, D = Z.shape
    # init
    U_0 = torch.randn(D, d_prev)
    U_0 = torch.linalg.qr(U_0)[0]

    # TRACKING VARS
    was_decreasing = False

    # init tracking variable
    # max_error = max_error_ratio*r_dimcheck

    # note kernel will stay the same for all d
    gamma = float(np.log(2)) / (r_dimcheck ** 2)
    kernel_pre = torch.exp(-gamma * (Z - z_c).pow(2).sum(dim=1, keepdim=True))
    z_mu = (Z * kernel_pre).sum(dim=0, keepdim=True) / kernel_pre.sum()
    kernel = torch.exp(-gamma * (Z - z_mu).pow(2).sum(dim=1, keepdim=True))

    # need to track U_prev in case of decreasing
    U_prev = U_0.clone()

    # note this loop iterates a maximum of D times
    for i in range(D):
        U, V, loss = opt_UV(Z, z_c, U_0, n_iter_inner, kernel=kernel)

        # compute max error, if too large, increase d
        coord = (Z - z_c) @ U
        Z_perp = (Z - z_c) - coord @ U.T
        # NOTE THIS NEEDS TO CHANGE IF DOING DIAG HESSIAN
        H_input = torch.bmm(coord.unsqueeze(2), coord.unsqueeze(1))
        # since tensor is symmetric, we only need to keep upper diag terms
        # output is a tensor of shape (N, d(d+1)/2)
        d = U.shape[1]
        idx_triu = torch.triu_indices(d, d)
        H_input = H_input[:, idx_triu[0, :], idx_triu[1, :]]
        l0_error = (kernel * (Z_perp - H_input @ V.T)).norm(dim=1).max()
        # l2_error = 0.5*(kernel * (Z_perp - H_input @ V.T)).norm(dim=1).pow(2).mean()
        # if achieved desired loss, reduce dimension
        # print(f'd: {d}, error: {l0_error}')
        if l0_error <= max_error:
            # if we were increasing but then achieved desired loss we have converged
            if (not was_decreasing and i > 0) or U.shape[1] == 1:
                break

            # otherwise, continue to reduce dimension
            was_decreasing = True
            # remove column of U lease correlated with V
            idx_min = ((Z - z_c) @ U).abs().mean(dim=0).argmin()
            U_0 = torch.cat((U[:, :idx_min], U[:, idx_min + 1:]), dim=1)
            U_prev = U.clone()

        else:
            # if we were decreasing but then achieved desired loss we have converged
            if was_decreasing or U.shape[1] == D:
                U = U_prev
                break
            U_new = torch.randn(D, 1)
            U_new = U_new - U @ (U.T @ U_new)
            U_new = U_new / torch.norm(U_new, p=2)
            U_0 = torch.cat((U, U_new), dim=1).clone()

    return U.detach().data, loss.detach().data
//  [python]
def __init__(self, d: int, x_min: float, x_max: float) ->None:
        super().__init__()
        self.x_min = x_min
        self.x_max = x_max
        M = torch.randn(d, d)
        Q, _ = torch.qr(M)
        self.register_buffer('Q', Q)
//  [python]
def __init__(self, d: int, x_min: float, x_max: float) ->None:
        super().__init__()
        self.x_min = x_min
        self.x_max = x_max
        M = torch.randn(d, d)
        Q, _ = torch.linalg.qr(M)
        self.register_buffer('Q', Q)
//  [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())
    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(), mode="reduced")
    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, channels=3, permutation=None):
        super(GeneralizedChannelPermute, self).__init__()
        self.__delattr__('permutation')
        W, _ = torch.qr(torch.randn(channels, channels))
        LU, pivots = torch.linalg.lu_factor(W)
        if permutation is None:
            P, _, _ = torch.lu_unpack(LU, pivots)
        else:
            if len(permutation) != channels:
                raise ValueError('Keyword argument "permutation" expected to have {} elements but {} found.'.format(channels, len(permutation)))
            P = torch.eye(channels, channels)[permutation.type(dtype=torch.int64)]
        self.register_buffer('permutation', P)
        self.LU = torch.nn.Parameter(LU)
//  [python]
def __init__(self, channels=3, permutation=None):
        super(GeneralizedChannelPermute, self).__init__()
        self.__delattr__('permutation')
        W, _ = torch.linalg.qr(torch.randn(channels, channels))
        LU, pivots = torch.linalg.lu_factor(W)
        if permutation is None:
            P, _, _ = torch.lu_unpack(LU, pivots)
        else:
            if len(permutation) != channels:
                raise ValueError('Keyword argument "permutation" expected to have {} elements but {} found.'.format(channels, len(permutation)))
            P = torch.eye(channels, channels)[permutation.type(dtype=torch.int64)]
        self.register_buffer('permutation', P)
        self.LU = torch.nn.Parameter(LU)
//  [python]
def transform_matrix(self, scale_matrix):
        # We require the scaling matrix to satisfy some constraints
        # such as being lower traingular with positive diagonal
        # entries (even when parameters are full).  This is similar to
        # requiring the standard deviation to be positive.  It's
        # transformed to an "equivalent" matrix satisfying the
        # necessary properties.

        # Unfortunately there's no QL factorization in PyTorch so we
        # transpose m and use the QR.  We only need the 'R' return
        # value, so the Q return value is dropped.

        upper = torch.qr(scale_matrix.T)[1]

        if not isinstance(self.univariate_model, UnivariateUnitScalingModel):
            # Normalize the colums (while in upper triangular form)
            # so that the univariate model determines the variances
            # and the multivariate model determines the correlations.
            upper = torch.nn.functional.normalize(upper, dim=0)

        # Transpose to make lower triangular, then make diagonal positive.
        lower = make_diagonal_nonnegative(upper.T)
        return lower
//  [python]
def transform_matrix(self, scale_matrix):
        # We require the scaling matrix to satisfy some constraints
        # such as being lower traingular with positive diagonal
        # entries (even when parameters are full).  This is similar to
        # requiring the standard deviation to be positive.  It's
        # transformed to an "equivalent" matrix satisfying the
        # necessary properties.

        # Unfortunately there's no QL factorization in PyTorch so we
        # transpose m and use the QR.  We only need the 'R' return
        # value, so the Q return value is dropped.

        upper = torch.linalg.qr(scale_matrix.T, mode="reduced")[1]

        if not isinstance(self.univariate_model, UnivariateUnitScalingModel):
            # Normalize the colums (while in upper triangular form)
            # so that the univariate model determines the variances
            # and the multivariate model determines the correlations.
            upper = torch.nn.functional.normalize(upper, dim=0)

        # Transpose to make lower triangular, then make diagonal positive.
        lower = make_diagonal_nonnegative(upper.T)
        return lower
//  [python]
def __init__(self, in_channels: int = 3):
        """Initializes weights.

        Args:
            in_channels: The number of channels of an input image.
        """

        super(InvConv2d, self).__init__()
        random_matrix = torch.randn(size=(in_channels, in_channels), dtype=torch.float32, device=self.device)
        w_init, _ = torch.qr(random_matrix)
        self.weight = nn.Parameter(rearrange(w_init, "h w -> h w () ()"))
//  [python]
def __init__(self, in_channels: int = 3):
        """Initializes weights.

        Args:
            in_channels: The number of channels of an input image.
        """

        super(InvConv2d, self).__init__()
        random_matrix = torch.randn(size=(in_channels, in_channels), dtype=torch.float32, device=self.device)
        w_init, _ = torch.linalg.qr(random_matrix)
        self.weight = nn.Parameter(rearrange(w_init, "h w -> h w () ()"))
//  [python]
def orthogonal_random_matrix_(w):
    """
    Modified slightly from https://github.com/idiap/fast-transformers/blob/master/fast_transformers/feature_maps/base.py
    to make orthogonal function conforming with PyTorch 1.9 linalg module.

    Initialize the matrix w in-place to compute orthogonal random features.

    The matrix is initialized such that its columns are orthogonal to each
    other (in groups of size `rows`) and their norms is drawn from the
    chi-square distribution with `rows` degrees of freedom (namely the norm of
    a `rows`-dimensional vector distributed as N(0, I)).

    Arguments
    ---------
        w: float tensor of size (rows, columns)
    """
    rows, columns = w.shape
    start = 0
    while start < columns:
        end = min(start+rows, columns)
        block = torch.randn(rows, rows, device=w.device)
        norms = torch.sqrt(torch.einsum("ab,ab->a", block, block))
        Q, _ = torch.qr(block)
        w[:, start:end] = (
            Q[:, :end-start] * norms[None, :end-start]
        )
        start += rows
//  [python]
def orthogonal_random_matrix_(w):
    """
    Modified slightly from https://github.com/idiap/fast-transformers/blob/master/fast_transformers/feature_maps/base.py
    to make orthogonal function conforming with PyTorch 1.9 linalg module.

    Initialize the matrix w in-place to compute orthogonal random features.

    The matrix is initialized such that its columns are orthogonal to each
    other (in groups of size `rows`) and their norms is drawn from the
    chi-square distribution with `rows` degrees of freedom (namely the norm of
    a `rows`-dimensional vector distributed as N(0, I)).

    Arguments
    ---------
        w: float tensor of size (rows, columns)
    """
    rows, columns = w.shape
    start = 0
    while start < columns:
        end = min(start+rows, columns)
        block = torch.randn(rows, rows, device=w.device)
        norms = torch.sqrt(torch.einsum("ab,ab->a", block, block))
        Q, _ = torch.linalg.qr(block)
        w[:, start:end] = (
            Q[:, :end-start] * norms[None, :end-start]
        )
        start += rows
//  [python]
def normal_space(m):
    d1, d2 = m.shape[-2:]
    assert d1 > d2
    return torch.qr(m)[0][..., -(d1-d2):]
//  [python]
def normal_space(m):
    d1, d2 = m.shape[-2:]
    assert d1 > d2
    return torch.linalg.qr(m, mode='complete')[0][..., -(d1-d2):]
//  [python]
def fake_picache_compress(
    input: torch.Tensor, loop, rank, device, p_base, q_base, quantize_bit, left
):
    input = input.float()
    p_base[0] = p_base[0].float()
    q_base[0] = q_base[0].float()
    batch, num_head, seq_len, sep_dim = input.shape
    input = (
        input.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )  # convert to 32bits for qr decomposition
    input = input.view(seq_len * batch, sep_dim * num_head)
    # 3 calculation = loop * (matmul) + 2 * qrO(n^2)
    for i in range(loop):
        if i == loop - 1:
            q, r = torch.qr(p_base[0])
            p_base[0] = q
        q_base[0] = input @ p_base[0]
        if i == loop - 1:
            q, r = torch.qr(q_base[0])
            q_base[0] = q
        p_base[0] = torch.transpose(input, 0, 1) @ q_base[0]
    # outlier quantization
    q_base[0] = fake_dense_sparse_uniformquantization(q_base[0], quantize_bit, left)
    p_base[0] = fake_dense_sparse_uniformquantization(p_base[0], quantize_bit, left)
    input = q_base[0] @ torch.transpose(p_base[0], 0, 1)
    input = input.view(batch, seq_len, num_head, sep_dim)
    input = input.permute(0, 2, 1, 3)
    input = input.half()
    p_base[0] = p_base[0].half()
    q_base[0] = q_base[0].half()
    return input
//  [python]
def fake_picache_compress(
    input: torch.Tensor, loop, rank, device, p_base, q_base, quantize_bit, left
):
    input = input.float()
    p_base[0] = p_base[0].float()
    q_base[0] = q_base[0].float()
    batch, num_head, seq_len, sep_dim = input.shape
    input = (
        input.permute(0, 2, 1, 3).contiguous().view(batch, seq_len, sep_dim * num_head)
    )  # convert to 32bits for qr decomposition
    input = input.view(seq_len * batch, sep_dim * num_head)
    # 3 calculation = loop * (matmul) + 2 * qrO(n^2)
    for i in range(loop):
        if i == loop - 1:
            p_base[0] = torch.linalg.qr(p_base[0]).Q
        q_base[0] = input @ p_base[0]
        if i == loop - 1:
            q_base[0] = torch.linalg.qr(q_base[0]).Q
        p_base[0] = torch.transpose(input, 0, 1) @ q_base[0]
    # outlier quantization
    q_base[0] = fake_dense_sparse_uniformquantization(q_base[0], quantize_bit, left)
    p_base[0] = fake_dense_sparse_uniformquantization(p_base[0], quantize_bit, left)
    input = q_base[0] @ torch.transpose(p_base[0], 0, 1)
    input = input.view(batch, seq_len, num_head, sep_dim)
    input = input.permute(0, 2, 1, 3)
    input = input.half()
    p_base[0] = p_base[0].half()
    q_base[0] = q_base[0].half()
    return input
//  [python]
def construct_approximation(
        A: lo.LinearOperator,
        l_0: int,
        l_max: int,
        power_iter_count: int,
        error_tol, dtype=None):
    n = A.shape[0]
    assert A.shape[0] == A.shape[1]
    device = A.device

    Y = torch.empty((n, 0), dtype=dtype, device=device)
    Omega = torch.empty((n, 0), dtype=dtype, device=device)
    E = torch.inf
    m = l_0
    break_early = False
    U = torch.Tensor([])
    Lambda_hat = torch.Tensor([])
    nu = math.sqrt(n) * torch.finfo().eps

    while E > error_tol:
        Omega_0 = torch.randn((n, m), dtype=dtype, device=device)
        Q, _ = torch.linalg.qr(Omega_0)
        Y_0 = lo.operator_matrix_product(A, Omega_0)
        Omega = torch.hstack([Omega, Omega_0])
        Y = torch.hstack([Y, Y_0])
        Y_nu = Y + nu * Omega
        C = torch.linalg.cholesky(Omega.T @ Y_nu).T
        B = torch.linalg.solve_triangular(C, Y_nu, upper=True, left=False)
        U, Sigma, _ = torch.linalg.svd(B, False)
        Lambda_hat = torch.relu(Sigma * Sigma - nu)

        if break_early:
            break

        E = randomized_power_err_est(A, U, Lambda_hat, power_iter_count, dtype, device)
        m, l_0 = l_0, 2 * l_0
        if l_0 > l_max:
            l_0 = l_0 - m
            m = l_max - l_0
            break_early = True
    #Add threshold to cut off too small values
    return _NystromApproximation(U, Lambda_hat)
//  [python]
def construct_approximation(
        A: lo.LinearOperator,
        l_0: int,
        l_max: int,
        power_iter_count: int,
        error_tol, dtype=None):
    n = A.shape[0]
    assert A.shape[0] == A.shape[1]
    device = A.device

    Y = torch.empty((n, 0), dtype=dtype, device=device)
    Omega = torch.empty((n, 0), dtype=dtype, device=device)
    E = torch.inf
    m = l_0
    break_early = False
    U = torch.Tensor([])
    Lambda_hat = torch.Tensor([])
    nu = math.sqrt(n) * torch.finfo().eps

    while E > error_tol:
        Omega_0 = torch.randn((n, m), dtype=dtype, device=device)
        Omega_0, _ = torch.linalg.qr(Omega_0)
        Y_0 = lo.operator_matrix_product(A, Omega_0)
        Omega = torch.hstack([Omega, Omega_0])
        Y = torch.hstack([Y, Y_0])
        Y_nu = Y + nu * Omega
        C = torch.linalg.cholesky(Omega.T @ Y_nu).T
        B = torch.linalg.solve_triangular(C, Y_nu, upper=True, left=False)
        U, Sigma, _ = torch.linalg.svd(B, False)
        Lambda_hat = torch.relu(Sigma * Sigma - nu)

        if break_early:
            break

        E = randomized_power_err_est(A, U, Lambda_hat, power_iter_count, dtype, device)
        m, l_0 = l_0, 2 * l_0
        if l_0 > l_max:
            l_0 = l_0 - m
            m = l_max - l_0
            break_early = True
    #Add threshold to cut off too small values
    return _NystromApproximation(U, Lambda_hat)
//  [python]
def forward(self, mat_feats):        
        W = self.isom_params.data
        
        q, r = torch.qr(W)        
        d = torch.diag(r, 0)
        ph = d.sign()
        q *= ph
        W.copy_(q)
        W = W.unsqueeze(0)
        
        mat_feats = W @ mat_feats @ W.transpose(-1, -2)
        
        return mat_feats
//  [python]
def forward(self, mat_feats):        
        W = self.isom_params.data
        
        q, r = torch.linalg.qr(W)        
        d = torch.diag(r, 0)
        ph = d.sign()
        q *= ph
        W.copy_(q)
        W = W.unsqueeze(0)
        
        mat_feats = W @ mat_feats @ W.transpose(-1, -2)
        
        return mat_feats
//  [python]
def _generate_W(F: int, k: int, device: str):
    omega = torch.randn((F, k), device=device)
    q, r = torch.qr(omega)
    W = q @ torch.sign(torch.diag(torch.diag(r)))
    return W
//  [python]
def _generate_W(F: int, k: int, device: str):
    omega = torch.randn((F, k), device=device)
    q, r = torch.linalg.qr(omega)
    W = q @ torch.sign(torch.diag(torch.diag(r)))
    return W
//  [python]
def __init__(self, input_size, hidden_size, num_layers=1, dropout=0.0):
        super(AffineRNN, self).__init__()
        self.num_layers = num_layers
        self.hidden_size = hidden_size
        self.dropout = dropout

        def rnn_weight_initialization(input_size, hidden_size):
            """
            the weight matrix is store as
            [   w_ih,
                w_hh,
                bias, ]
            
            The normal_ initialization might be better than uniform_ (keep the same variance)? 
            """
            w = torch.empty(input_size + hidden_size, hidden_size).normal_(
                std=(3*(input_size + hidden_size)) ** (-0.5)
            )
            # w = torch.empty(input_size + hidden_size, hidden_size).uniform_(
            #     -(input_size+hidden_size)**(-0.5), (input_size+hidden_size)**(-0.5)
            # )
            w[input_size:] = torch.qr(w[input_size:])[0]
            b = torch.zeros(1, hidden_size)
            return torch.cat([w, b])

        self.param0 = torch.nn.Parameter(rnn_weight_initialization(input_size, hidden_size))
        for i in range(num_layers - 1):
            self.register_parameter(
                f"param{i+1}",
                torch.nn.Parameter(rnn_weight_initialization(hidden_size, hidden_size)),
            )
//  [python]
def __init__(self, input_size, hidden_size, num_layers=1, dropout=0.0):
        super(AffineRNN, self).__init__()
        self.num_layers = num_layers
        self.hidden_size = hidden_size
        self.dropout = dropout

        def rnn_weight_initialization(input_size, hidden_size):
            """
            the weight matrix is store as
            [   w_ih,
                w_hh,
                bias, ]
            
            The normal_ initialization might be better than uniform_ (keep the same variance)? 
            """
            w = torch.empty(input_size + hidden_size, hidden_size).normal_(
                std=(3*(input_size + hidden_size)) ** (-0.5)
            )
            # w = torch.empty(input_size + hidden_size, hidden_size).uniform_(
            #     -(input_size+hidden_size)**(-0.5), (input_size+hidden_size)**(-0.5)
            # )
            w[input_size:] = torch.linalg.qr(w[input_size:])[0]
            b = torch.zeros(1, hidden_size)
            return torch.cat([w, b])

        self.param0 = torch.nn.Parameter(rnn_weight_initialization(input_size, hidden_size))
        for i in range(num_layers - 1):
            self.register_parameter(
                f"param{i+1}",
                torch.nn.Parameter(rnn_weight_initialization(hidden_size, hidden_size)),
            )
//  [python]
def sliced_wasserstein_sphere_uniform(Xs, num_projections):
    """
    Compute the sliced-Wasserstein distance on the sphere with uniform distrib.

    """
    d = Xs.shape[1]

    ## Uniforms and independent samples on the Stiefel manifold V_{d,2}
    Z = torch.randn((num_projections,d,2), device=Xs.device)
    Q, _ = torch.qr(Z)

    return sliced_cost_unif(Xs, Q)
//  [python]
def sliced_wasserstein_sphere_uniform(Xs, num_projections):
    """
    Compute the sliced-Wasserstein distance on the sphere with uniform distrib.

    """
    d = Xs.shape[1]

    ## Uniforms and independent samples on the Stiefel manifold V_{d,2}
    Z = torch.randn((num_projections,d,2), device=Xs.device)
    U, _ = torch.linalg.qr(Z)

    return sliced_cost_unif(Xs, U)
//  [python]
def matrix_from_singular_values(self, cond):
        log_cond = np.log(self.cond)
        log_s = torch.arange(-log_cond / 4., log_cond * (self.d + 1) / (4 * (self.d - 1)),
                             log_cond / (2. * (self.d - 1)))
        log_s = log_s.repeat(self.N, 1)
        log_s = log_s[:, :, None]
        s = torch.exp(log_s)
        S = torch.eye(s.shape[1]) * s[:, None]
        U, _ = torch.qr((torch.rand(self.N, 1, self.d, self.d) - 5.) * 200)
        V, _ = torch.qr((torch.rand(self.N, 1, self.d, self.d) - 5.) * 200)
        X = torch.matmul(U, torch.matmul(S, V.transpose(2, 3)))
        X = torch.matmul(X, X.transpose(2, 3))
        self.X = X
//  [python]
def matrix_from_singular_values(self, cond):
        log_cond = np.log(self.cond)
        log_s = torch.arange(-log_cond / 4., log_cond * (self.d + 1) / (4 * (self.d - 1)),
                             log_cond / (2. * (self.d - 1)))
        log_s = log_s.repeat(self.N, 1)
        log_s = log_s[:, :, None]
        s = torch.exp(log_s)
        S = torch.eye(s.shape[1]) * s[:, None]
        U, _ = torch.linalg.qr((torch.rand(self.N, 1, self.d, self.d) - 5.) * 200)
        V, _ = torch.linalg.qr((torch.rand(self.N, 1, self.d, self.d) - 5.) * 200)
        X = torch.matmul(U, torch.matmul(S, V.transpose(2, 3)))
        X = torch.matmul(X, X.transpose(2, 3))
        self.X = X
//  [python]
def sliced_wasserstein_sphere(Xs, Xt, num_projections, u_weights=None, v_weights=None, p=2, method=0):
    """
        Compute the sliced-Wasserstein distance on the sphere.

        Parameters:
        Xs: ndarray, shape (n_samples_u, dim)
            Samples in the source domain
        Xt: ndarray, shape (n_samples_v, dim)
            Samples in the target domain
        num_projections: int
            Number of projections
        device: str
        p: float
            Power of SW. Need to be >= 1.
        method: int
            0 for binary search, 1 for level median.
    """
    d = Xs.size(1)

    ## Uniforms and independent samples on the Stiefel manifold V_{d,2}
    Z = torch.randn((num_projections,d,2), device=Xs.device)
    U, _ = torch.qr(Z)

    return sliced_cost(Xs, Xt, U, p=p, u_weights=u_weights, v_weights=v_weights, method=method)
//  [python]
def sliced_wasserstein_sphere(Xs, Xt, num_projections, u_weights=None, v_weights=None, p=2, method=0):
    """
        Compute the sliced-Wasserstein distance on the sphere.

        Parameters:
        Xs: ndarray, shape (n_samples_u, dim)
            Samples in the source domain
        Xt: ndarray, shape (n_samples_v, dim)
            Samples in the target domain
        num_projections: int
            Number of projections
        device: str
        p: float
            Power of SW. Need to be >= 1.
        method: int
            0 for binary search, 1 for level median.
    """
    d = Xs.size(1)

    ## Uniforms and independent samples on the Stiefel manifold V_{d,2}
    Z = torch.randn((num_projections,d,2), device=Xs.device)
    U, _ = torch.linalg.qr(Z)

    return sliced_cost(Xs, Xt, U, p=p, u_weights=u_weights, v_weights=v_weights, method=method)
//  [python]
def __init__(self, output_size, train_values, all_values, scaler="log"):
        super(DICE, self).__init__(output_size,train_values, all_values)

        self.M = torch.randn(self.output_size, self.output_size)
        self.Q, self.R = torch.qr(self.M)

        if torch.cuda.is_available():
            self.Q = self.Q.cuda()
            self.R = self.R.cuda()

        self.scaler_name  = scaler

        if scaler == "log":
            self.scaler = log_scaling_function
        elif scaler == "quantile":
            self.scaler = QuantileScaling(self.train_values).quantile_scaling_function

       
        
        train_value_list = sorted(self.train_values)
        
        tensor_train_values = self.scaler(torch.tensor(train_value_list))

        

        self.a = tensor_train_values[0]
        self.b = tensor_train_values[-1]


        if torch.cuda.is_available() :
            self.dice_embedding = torch.nn.Embedding(len(self.all_values), output_size).cuda()
        else:
            self.dice_embedding = torch.nn.Embedding(len(self.all_values), output_size)
        self.dice_embedding.weight.requires_grad = False


        all_values_tensor =torch.tensor(self.all_values)
        if torch.cuda.is_available() :
            all_values_tensor = all_values_tensor.cuda()

        
        x = self.scaler(all_values_tensor)

        theta = (x -self.a) / (self.b -self.a)  * np.pi

        
        # self.dice_embedding.weight = self.encode(all_values_tensor)

        sin_theta = torch.sin(theta)
        cos_theta = torch.cos(theta)


        d = torch.arange(0, self.output_size -1 ).float().to(x.device)

        self.dice_embedding.weight[:, :-1] = torch.pow(sin_theta.unsqueeze(1), d.unsqueeze(0)) * cos_theta.unsqueeze(1) 
        self.dice_embedding.weight[:, -1] = torch.pow(sin_theta, self.output_size)  


        self.dice_embedding.weight = nn.Parameter(torch.mm(self.dice_embedding.weight, self.Q))
        self.dice_embedding.weight.requires_grad = False
//  [python]
def __init__(self, output_size, train_values, all_values, scaler="log"):
        super(DICE, self).__init__(output_size,train_values, all_values)

        self.M = torch.randn(self.output_size, self.output_size)
        self.Q, self.R = torch.linalg.qr(self.M)

        if torch.cuda.is_available():
            self.Q = self.Q.cuda()
            self.R = self.R.cuda()

        self.scaler_name  = scaler

        if scaler == "log":
            self.scaler = log_scaling_function
        elif scaler == "quantile":
            self.scaler = QuantileScaling(self.train_values).quantile_scaling_function

       
        
        train_value_list = sorted(self.train_values)
        
        tensor_train_values = self.scaler(torch.tensor(train_value_list))

        

        self.a = tensor_train_values[0]
        self.b = tensor_train_values[-1]


        if torch.cuda.is_available() :
            self.dice_embedding = torch.nn.Embedding(len(self.all_values), output_size).cuda()
        else:
            self.dice_embedding = torch.nn.Embedding(len(self.all_values), output_size)
        self.dice_embedding.weight.requires_grad = False


        all_values_tensor =torch.tensor(self.all_values)
        if torch.cuda.is_available() :
            all_values_tensor = all_values_tensor.cuda()

        
        x = self.scaler(all_values_tensor)

        theta = (x -self.a) / (self.b -self.a)  * np.pi

        
        # self.dice_embedding.weight = self.encode(all_values_tensor)

        sin_theta = torch.sin(theta)
        cos_theta = torch.cos(theta)


        d = torch.arange(0, self.output_size -1 ).float().to(x.device)

        self.dice_embedding.weight[:, :-1] = torch.pow(sin_theta.unsqueeze(1), d.unsqueeze(0)) * cos_theta.unsqueeze(1) 
        self.dice_embedding.weight[:, -1] = torch.pow(sin_theta, self.output_size)  


        self.dice_embedding.weight = nn.Parameter(torch.mm(self.dice_embedding.weight, self.Q))
        self.dice_embedding.weight.requires_grad = False
//  [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 forward(a: Array) -> Any:
        q, r = torch.qr(a)
        # ctx.save_for_backward(a, q, r)
        return q, r
//  [python]
def forward(a: Array) -> Any:
        q, r = torch.linalg.qr(a, mode="reduced")
        # ctx.save_for_backward(a, q, r)
        return q, r