import torch
import flashinfer

torch.manual_seed(42)

batch_size = 128
hidden_dim = 4096
out = torch.empty(batch_size, hidden_dim // 2, dtype=torch.float16)
input_tensor = torch.randn(batch_size, hidden_dim, dtype=torch.float16)
enable_pdl = True

# make sure the output tensor is on the GPU if available
if torch.cuda.is_available():
    out = out.cuda()
    input_tensor = input_tensor.cuda()
else:
    raise RuntimeError("CUDA is not available. Please run this on a GPU.")

flashinfer.gelu_and_mul(out, input_tensor, enable_pdl)
print("GELU and multiply operation completed. Output shape:", out.shape)
print("Output tensor sample:", out[:, :5])  # Show first 5 elements