Function f

examples/torch_cuda_kernel.py:11–11

(tg_out, tg_data)

@TinyJit
def f(tg_out, tg_data): return tg_out.assign(tg_data[:, :, 0] * 0.2989 + tg_data[:, :, 1] * 0.5870 + tg_data[:, :, 2] * 0.1140).realize()

def custom_kernel(data: torch.Tensor, device="CUDA") -> torch.Tensor:
  assert data.dtype == torch.float32