class_docstr: Tests just-in-time cpp extensions.
  Don't confuse this with the PyTorch JIT (aka TorchScript).
code: def test_inline_jit_compile_custom_op_cuda(self):
      cuda_source = """
          __global__ void cos_add_kernel(
              const float* __restrict__ x,
              const float* __restrict__ y,
              float* __restrict__ output,
              const int size) {
            const auto index = blockIdx.x * blockDim.x + threadIdx.x;
            if (index < size) {
              output[index] = __cosf(x[index]) + __cosf(y[index]);
            }
          }
  
          torch::Tensor cos_add(torch::Tensor x, torch::Tensor y) {
            auto output = torch::zeros_like(x);
            const int threads = 1024;
            const int blocks = (output.numel() + threads - 1) / threads;
            cos_add_kernel<<<blocks, threads>>>(x.data_ptr<float>(), y.data_ptr<float>(), output.data_ptr<float>(), output.numel());
            return output;
          }
          """
  
      # Here, the C++ source need only declare the function signature.
      cpp_source = """
             #include <torch/library.h>
             torch::Tensor cos_add(torch::Tensor x, torch::Tensor y);
  
             TORCH_LIBRARY(inline_jit_extension_custom_op_cuda, m) {
               m.def("cos_add", cos_add);
             }
          """
  
      torch.utils.cpp_extension.load_inline(
          name="inline_jit_extension_custom_op_cuda",
          cpp_sources=cpp_source,
          cuda_sources=cuda_source,
          verbose=True,
          is_python_module=False,
      )
  
      x = torch.randn(4, 4, device="cuda", dtype=torch.float32)
      y = torch.randn(4, 4, device="cuda", dtype=torch.float32)
  
      z = torch.ops.inline_jit_extension_custom_op_cuda.cos_add(x, y)
      self.assertEqual(z, x.cos() + y.cos())
  
imports: import os
  import shutil
  import sys
  import unittest
  import warnings
  import re
  import tempfile
  import subprocess
  import glob
  import torch.testing._internal.common_utils
  from torch.testing._internal.common_cuda import TEST_CUDNN, TEST_CUDA
  import torch
  import torch.backends.cudnn
  import torch.utils.cpp_extension
  from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
  from torch.testing._internal.common_utils import gradcheck
  import torch.multiprocessing
  from torch.utils.cpp_extension import _TORCH_PATH, remove_extension_h_precompiler_headers, get_cxx_compiler, check_compiler_is_gcc
label_desc: Use image node_id 15 for calling the TestCppExtensionJIT obj's underlying member method code with example usage: obj.test_inline_jit_compile_custom_op_cuda() without return types