# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "torch",
#     "triton",
#     "kernels",
# ]
# ///
"""Minimal smoke test of the published kernel via get_kernel (run on a CUDA box)."""

import torch
from kernels import get_kernel


kir = get_kernel("Molbap/kernel_image_resize", revision="main", trust_remote_code=True)

device = "cuda" if torch.cuda.is_available() else "cpu"
images = [
    torch.randint(0, 256, (3, h, w), dtype=torch.uint8, device=device)
    for h, w in [(640, 480), (800, 600), (384, 1024)]
]

pixel_values = kir.resize_normalize(
    images,
    size=384,
    image_mean=[0.5, 0.5, 0.5],
    image_std=[0.5, 0.5, 0.5],
    rescale_factor=1 / 255,
    resample="bicubic",
    antialias=True,
)
print(f"{len(images)} ragged images -> {tuple(pixel_values.shape)} {pixel_values.dtype}")