import torch
from deepinfer.model_executor.layers.tilelang_ops.neuron.random_sampling import philox4x32_10
def t(v): return torch.tensor([v], dtype=torch.int64)
for ctr, key, exp in [((0,0,0,0),(0,0),(0x6627e8d5,0xe169c58d,0xbc57ac4c,0x9b00dbd8)),
    ((0xffffffff,)*4,(0xffffffff,)*2,(0x408f276d,0x41c83b0e,0xa20bc7c6,0x6d5451fd)),
    ((0x243f6a88,0x85a308d3,0x13198a2e,0x03707344),(0xa4093822,0x299f31d0),(0xd16cfe09,0x94fdcceb,0x5001e420,0x24126ea1))]:
    out = philox4x32_10(*map(t, ctr), *map(t, key))
    got = tuple(int(o) for o in out)
    print([hex(g) for g in got], got == exp)
from deepinfer.model_executor.layers.tilelang_ops import ops
print(ops.random_sampling_tl, ops.fused_sample_tl, ops.map_lookup_tl)