from aiter.ops.shuffle import ck_shuffle_weight
import torch
import numpy as np
torch.set_printoptions(profile="full")
input = torch.arange(786432).cuda()
input = input.view(3,256,1024)
print(input)
output = ck_shuffle_weight(input)
# for i in range(output.shape[0]):
#     for j in range(output.shape[1]):
#         np.savetxt(f'tensor_slice_{i}_{j}.txt', output[i][j].cpu().numpy(), fmt='%d')
print(output)
print(output.shape)