text: import os
  import shutil
  import sys
  import unittest
  import warnings
  import re
  import tempfile
  import subprocess
  import glob
  
  import textwrap
  from multiprocessing import Process
  
  import torch.testing._internal.common_utils as common
  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, TEST_WITH_ASAN, has_breakpad
  
  
  TEST_CUDA = torch.cuda.is_available() and CUDA_HOME is not None
  TEST_CUDNN = False
  TEST_ROCM = torch.cuda.is_available() and torch.version.hip is not None and ROCM_HOME is not None
  if TEST_CUDA and torch.version.cuda is not None:  # the skip CUDNN test for ROCm
      CUDNN_HEADER_EXISTS = os.path.isfile(os.path.join(CUDA_HOME, "include/cudnn.h"))
      TEST_CUDNN = (
          TEST_CUDA and CUDNN_HEADER_EXISTS and torch.backends.cudnn.is_available()
      )
  IS_WINDOWS = sys.platform == "win32"
  
  
  def remove_build_path():
      if sys.platform == "win32":
          print("Not wiping extensions build folder because Windows")
          return
      default_build_root = torch.utils.cpp_extension.get_default_build_root()
      if os.path.exists(default_build_root):
          shutil.rmtree(default_build_root)
  
  
  class TestCppExtensionJIT(common.TestCase):
      """Tests just-in-time cpp extensions.
      Don't confuse this with the PyTorch JIT (aka TorchScript).
      """
  
      def setUp(self):
          super().setUp()
          # cpp extensions use relative paths. Those paths are relative to
          # this file, so we'll change the working directory temporarily
          self.old_working_dir = os.getcwd()
          os.chdir(os.path.dirname(os.path.abspath(__file__)))
  
      def tearDown(self):
          super().tearDown()
          # return the working directory (see setUp)
          os.chdir(self.old_working_dir)
  
      @classmethod
      def setUpClass(cls):
          remove_build_path()
  
      @classmethod
      def tearDownClass(cls):
          remove_build_path()
  
      def test_jit_compile_extension(self):
          module = torch.utils.cpp_extension.load(
              name="jit_extension",
              sources=[
                  "cpp_extensions/jit_extension.cpp",
                  "cpp_extensions/jit_extension2.cpp",
              ],
              extra_include_paths=["cpp_extensions"],
              extra_cflags=["-g"],
              verbose=True,
          )
          x = torch.randn(4, 4)
          y = torch.randn(4, 4)
  
          z = module.tanh_add(x, y)
          self.assertEqual(z, x.tanh() + y.tanh())
  
          # Checking we can call a method defined not in the main C++ file.
          z = module.exp_add(x, y)
          self.assertEqual(z, x.exp() + y.exp())
  
          # Checking we can use this JIT-compiled class.
          doubler = module.Doubler(2, 2)
          self.assertIsNone(doubler.get().grad)
          self.assertEqual(doubler.get().sum(), 4)
          self.assertEqual(doubler.forward().sum(), 8)
  
      @unittest.skipIf(not (TEST_CUDA or TEST_ROCM), "CUDA not found")
      def test_jit_cuda_extension(self):
          # NOTE: The name of the extension must equal the name of the module.
          module = torch.utils.cpp_extension.load(
              name="torch_test_cuda_extension",
              sources=[
                  "cpp_extensions/cuda_extension.cpp",
                  "cpp_extensions/cuda_extension.cu",
              ],
              extra_cuda_cflags=["-O2"],
              verbose=True,
              keep_intermediates=False,
          )
  
          x = torch.zeros(100, device="cuda", dtype=torch.float32)
          y = torch.zeros(100, device="cuda", dtype=torch.float32)
  
          z = module.sigmoid_add(x, y).cpu()
  
          # 2 * sigmoid(0) = 2 * 0.5 = 1
          self.assertEqual(z, torch.ones_like(z))
  
      def _run_jit_cuda_archflags(self, flags, expected):
          # Compile an extension with given `flags`
          def _check_cuobjdump_output(expected_values, is_ptx=False):
              elf_or_ptx = '--list-ptx' if is_ptx else '--list-elf'
              lib_ext = '.pyd' if IS_WINDOWS else '.so'
              # Note, .extension name may include _v1, _v2, so first find exact name
              ext_filename = glob.glob(os.path.join(temp_dir,
                                                    'cudaext_archflag*' + lib_ext))[0]
              command = ['cuobjdump', elf_or_ptx, ext_filename]
              p = subprocess.Popen(command,
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE)
              output, err = p.communicate()
              output = output.decode("ascii")
              err = err.decode("ascii")
  
              if not p.returncode == 0 or not err == '':
                  raise AssertionError("Flags: {}\nReturncode: {}\nStderr: {}\n"
                                       "Output: {} ".format(flags, p.returncode,
                                                            err, output))
  
              actual_arches = sorted(re.findall(r'sm_\d\d', output))
              expected_arches = sorted(['sm_' + xx for xx in expected_values])
              self.assertEqual(actual_arches, expected_arches,
                               msg="Flags: {},  Actual: {},  Expected: {}\n"
                                   "Stderr: {}\nOutput: {}".format(
                                       flags, actual_arches, expected_arches,
                                       err, output))
  
          temp_dir = tempfile.mkdtemp()
          old_envvar = os.environ.get('TORCH_CUDA_ARCH_LIST', None)
          try:
              os.environ['TORCH_CUDA_ARCH_LIST'] = flags
              torch.utils.cpp_extension.load(
                  name="cudaext_archflags",
                  sources=[
                      "cpp_extensions/cuda_extension.cpp",
                      "cpp_extensions/cuda_extension.cu",
                  ],
                  extra_cuda_cflags=["-O2"],
                  verbose=True,
                  build_directory=temp_dir,
              )
  
              # Expected output for --list-elf:
              #   ELF file    1: cudaext_archflags.1.sm_61.cubin
              #   ELF file    2: cudaext_archflags.2.sm_52.cubin
              _check_cuobjdump_output(expected[0])
              if expected[1] is not None:
                  # Expected output for --list-ptx:
                  #   PTX file    1: cudaext_archflags.1.sm_61.ptx
                  _check_cuobjdump_output(expected[1], is_ptx=True)
          finally:
              if IS_WINDOWS:
                  print("Not wiping extensions build folder because Windows")
              else:
                  shutil.rmtree(temp_dir)
  
              if old_envvar is None:
                  os.environ.pop('TORCH_CUDA_ARCH_LIST')
              else:
                  os.environ['TORCH_CUDA_ARCH_LIST'] = old_envvar
  
      @unittest.skipIf(not TEST_CUDA, "CUDA not found")
      @unittest.skipIf(TEST_ROCM, "disabled on rocm")
      def test_jit_cuda_archflags(self):
          # Test a number of combinations:
          #   - the default for the machine we're testing on
          #   - Separators, can be ';' (most common) or ' '
          #   - Architecture names
          #   - With/without '+PTX'
  
          n = torch.cuda.device_count()
          capabilities = {torch.cuda.get_device_capability(i) for i in range(n)}
          # expected values is length-2 tuple: (list of ELF, list of PTX)
          # note: there should not be more than one PTX value
          archflags = {
              '': (['{}{}'.format(capability[0], capability[1]) for capability in capabilities], None),
              "Maxwell+Tegra;6.1": (['53', '61'], None),
              "Pascal 3.5": (['35', '60', '61'], None),
              "Volta": (['70'], ['70']),
          }
          if int(torch.version.cuda.split('.')[0]) >= 10:
              # CUDA 9 only supports compute capability <= 7.2
              archflags["7.5+PTX"] = (['75'], ['75'])
              archflags["5.0;6.0+PTX;7.0;7.5"] = (['50', '60', '70', '75'], ['60'])
  
          for flags, expected in archflags.items():
              self._run_jit_cuda_archflags(flags, expected)
  
      @unittest.skipIf(not TEST_CUDNN, "CuDNN not found")
      def test_jit_cudnn_extension(self):
          # implementation of CuDNN ReLU
          if IS_WINDOWS:
              extra_ldflags = ["cudnn.lib"]
          else:
              extra_ldflags = ["-lcudnn"]
          module = torch.utils.cpp_extension.load(
              name="torch_test_cudnn_extension",
              sources=["cpp_extensions/cudnn_extension.cpp"],
              extra_ldflags=extra_ldflags,
              verbose=True,
              with_cuda=True,
          )
  
          x = torch.randn(100, device="cuda", dtype=torch.float32)
          y = torch.zeros(100, device="cuda", dtype=torch.float32)
          module.cudnn_relu(x, y)  # y=relu(x)
          self.assertEqual(torch.nn.functional.relu(x), y)
          with self.assertRaisesRegex(RuntimeError, "same size"):
              y_incorrect = torch.zeros(20, device="cuda", dtype=torch.float32)
              module.cudnn_relu(x, y_incorrect)
  
      def test_inline_jit_compile_extension_with_functions_as_list(self):
          cpp_source = """
          torch::Tensor tanh_add(torch::Tensor x, torch::Tensor y) {
            return x.tanh() + y.tanh();
          }
          """
  
          module = torch.utils.cpp_extension.load_inline(
              name="inline_jit_extension_with_functions_list",
              cpp_sources=cpp_source,
              functions="tanh_add",
              verbose=True,
          )
  
          self.assertEqual(module.tanh_add.__doc__.split("\n")[2], "tanh_add")
  
          x = torch.randn(4, 4)
          y = torch.randn(4, 4)
  
          z = module.tanh_add(x, y)
          self.assertEqual(z, x.tanh() + y.tanh())
  
      def test_inline_jit_compile_extension_with_functions_as_dict(self):
          cpp_source = """
          torch::Tensor tanh_add(torch::Tensor x, torch::Tensor y) {
            return x.tanh() + y.tanh();
          }
          """
  
          module = torch.utils.cpp_extension.load_inline(
              name="inline_jit_extension_with_functions_dict",
              cpp_sources=cpp_source,
              functions={"tanh_add": "Tanh and then sum :D"},
              verbose=True,
          )
  
          self.assertEqual(module.tanh_add.__doc__.split("\n")[2], "Tanh and then sum :D")
  
      def test_inline_jit_compile_extension_multiple_sources_and_no_functions(self):
          cpp_source1 = """
          torch::Tensor sin_add(torch::Tensor x, torch::Tensor y) {
            return x.sin() + y.sin();
          }
          """
  
          cpp_source2 = """
          #include <torch/extension.h>
          torch::Tensor sin_add(torch::Tensor x, torch::Tensor y);
          PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
            m.def("sin_add", &sin_add, "sin(x) + sin(y)");
          }
          """
  
          module = torch.utils.cpp_extension.load_inline(
              name="inline_jit_extension",
              cpp_sources=[cpp_source1, cpp_source2],
              verbose=True,
          )
  
          x = torch.randn(4, 4)
          y = torch.randn(4, 4)
  
          z = module.sin_add(x, y)
          self.assertEqual(z, x.sin() + y.sin())
  
      @unittest.skip("Temporarily disabled")
      @unittest.skipIf(not (TEST_CUDA or TEST_ROCM), "CUDA not found")
      def test_inline_jit_compile_extension_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<float>(), y.data<float>(), output.data<float>(), output.numel());
            return output;
          }
          """
  
          # Here, the C++ source need only declare the function signature.
          cpp_source = "torch::Tensor cos_add(torch::Tensor x, torch::Tensor y);"
  
          module = torch.utils.cpp_extension.load_inline(
              name="inline_jit_extension_cuda",
              cpp_sources=cpp_source,
              cuda_sources=cuda_source,
              functions=["cos_add"],
              verbose=True,
          )
  
          self.assertEqual(module.cos_add.__doc__.split("\n")[2], "cos_add")
  
          x = torch.randn(4, 4, device="cuda", dtype=torch.float32)
          y = torch.randn(4, 4, device="cuda", dtype=torch.float32)
  
          z = module.cos_add(x, y)
          self.assertEqual(z, x.cos() + y.cos())
  
      @unittest.skip("Temporarily disabled")
      @unittest.skipIf(not (TEST_CUDA or TEST_ROCM), "CUDA not found")
      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())
  
      def test_inline_jit_compile_extension_throws_when_functions_is_bad(self):
          with self.assertRaises(ValueError):
              torch.utils.cpp_extension.load_inline(
                  name="invalid_jit_extension", cpp_sources="", functions=5
              )
  
      def test_lenient_flag_handling_in_jit_extensions(self):
          cpp_source = """
          torch::Tensor tanh_add(torch::Tensor x, torch::Tensor y) {
            return x.tanh() + y.tanh();
          }
          """
  
          module = torch.utils.cpp_extension.load_inline(
              name="lenient_flag_handling_extension",
              cpp_sources=cpp_source,
              functions="tanh_add",
              extra_cflags=["-g\n\n", "-O0 -Wall"],
              extra_include_paths=["       cpp_extensions\n"],
              verbose=True,
          )
  
          x = torch.zeros(100, dtype=torch.float32)
          y = torch.zeros(100, dtype=torch.float32)
          z = module.tanh_add(x, y).cpu()
          self.assertEqual(z, x.tanh() + y.tanh())
  
      @unittest.skip("Temporarily disabled")
      @unittest.skipIf(not (TEST_CUDA or TEST_ROCM), "CUDA not found")
      def test_half_support(self):
          """
          Checks for an issue with operator< ambiguity for half when certain
          THC headers are included.
  
          See https://github.com/pytorch/pytorch/pull/10301#issuecomment-416773333
          for the corresponding issue.
          """
          cuda_source = """
          #include <THC/THCGeneral.h>
  
          template<typename T, typename U>
          __global__ void half_test_kernel(const T* input, U* output) {
              if (input[0] < input[1] || input[0] >= input[1]) {
                  output[0] = 123;
              }
          }
  
          torch::Tensor half_test(torch::Tensor input) {
              auto output = torch::empty(1, input.options().dtype(torch::kFloat));
              AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.scalar_type(), "half_test", [&] {
                  half_test_kernel<scalar_t><<<1, 1>>>(
                      input.data<scalar_t>(),
                      output.data<float>());
              });
              return output;
          }
          """
  
          module = torch.utils.cpp_extension.load_inline(
              name="half_test_extension",
              cpp_sources="torch::Tensor half_test(torch::Tensor input);",
              cuda_sources=cuda_source,
              functions=["half_test"],
              verbose=True,
          )
  
          x = torch.randn(3, device="cuda", dtype=torch.half)
          result = module.half_test(x)
          self.assertEqual(result[0], 123)
  
      def test_reload_jit_extension(self):
          def compile(code):
              return torch.utils.cpp_extension.load_inline(
                  name="reloaded_jit_extension",
                  cpp_sources=code,
                  functions="f",
                  verbose=True,
              )
  
          module = compile("int f() { return 123; }")
          self.assertEqual(module.f(), 123)
  
          module = compile("int f() { return 456; }")
          self.assertEqual(module.f(), 456)
          module = compile("int f() { return 456; }")
          self.assertEqual(module.f(), 456)
  
          module = compile("int f() { return 789; }")
          self.assertEqual(module.f(), 789)
  
      def test_cpp_frontend_module_has_same_output_as_python(self, dtype=torch.double):
          extension = torch.utils.cpp_extension.load(
              name="cpp_frontend_extension",
              sources="cpp_extensions/cpp_frontend_extension.cpp",
              verbose=True,
          )
  
          input = torch.randn(2, 5, dtype=dtype)
          cpp_linear = extension.Net(5, 2)
          cpp_linear.to(dtype)
          python_linear = torch.nn.Linear(5, 2).to(dtype)
  
          # First make sure they have the same parameters
          cpp_parameters = dict(cpp_linear.named_parameters())
          with torch.no_grad():
              python_linear.weight.copy_(cpp_parameters["fc.weight"])
              python_linear.bias.copy_(cpp_parameters["fc.bias"])
  
          cpp_output = cpp_linear.forward(input)
          python_output = python_linear(input)
          self.assertEqual(cpp_output, python_output)
  
          cpp_output.sum().backward()
          python_output.sum().backward()
  
          for p in cpp_linear.parameters():
              self.assertFalse(p.grad is None)
  
          self.assertEqual(cpp_parameters["fc.weight"].grad, python_linear.weight.grad)
          self.assertEqual(cpp_parameters["fc.bias"].grad, python_linear.bias.grad)
  
      def test_cpp_frontend_module_python_inter_op(self):
          extension = torch.utils.cpp_extension.load(
              name="cpp_frontend_extension",
              sources="cpp_extensions/cpp_frontend_extension.cpp",
              verbose=True,
          )
  
          # Create a torch.nn.Module which uses the C++ module as a submodule.
          class M(torch.nn.Module):
              def __init__(self):
                  super(M, self).__init__()
                  self.x = torch.nn.Parameter(torch.tensor(1.0))
                  self.net = extension.Net(3, 5)
  
              def forward(self, input):
                  return self.net.forward(input) + self.x
  
          net = extension.Net(5, 2)
          net.double()
          net.to(torch.get_default_dtype())
          self.assertEqual(str(net), "Net")
  
          # Further embed the torch.nn.Module into a Sequential, and also add the
          # C++ module as an element of the Sequential.
          sequential = torch.nn.Sequential(M(), torch.nn.Tanh(), net, torch.nn.Sigmoid())
  
          input = torch.randn(2, 3)
          # Try calling the module!
          output = sequential.forward(input)
          # The call operator is bound to forward too.
          self.assertEqual(output, sequential(input))
          self.assertEqual(list(output.shape), [2, 2])
  
          # Do changes on the module hierarchy.
          old_dtype = torch.get_default_dtype()
          sequential.to(torch.float64)
          sequential.to(torch.float32)
          sequential.to(old_dtype)
          self.assertEqual(sequential[2].parameters()[0].dtype, old_dtype)
  
          # Make sure we can access these methods recursively.
          self.assertEqual(len(list(sequential.parameters())), len(net.parameters()) * 2 + 1)
          self.assertEqual(len(list(sequential.named_parameters())), len(net.named_parameters()) * 2 + 1)
          self.assertEqual(len(list(sequential.buffers())), len(net.buffers()) * 2)
          self.assertEqual(len(list(sequential.modules())), 8)
  
          # Test clone()
          net2 = net.clone()
          self.assertEqual(len(net.parameters()), len(net2.parameters()))
          self.assertEqual(len(net.buffers()), len(net2.buffers()))
          self.assertEqual(len(net.modules()), len(net2.modules()))
  
          # Try differentiating through the whole module.
          for parameter in net.parameters():
              self.assertIsNone(parameter.grad)
          output.sum().backward()
          for parameter in net.parameters():
              self.assertFalse(parameter.grad is None)
              self.assertGreater(parameter.grad.sum(), 0)
  
          # Try calling zero_grad()
          net.zero_grad()
          for p in net.parameters():
              self.assertEqual(p.grad, torch.zeros_like(p))
  
          # Test train(), eval(), training (a property)
          self.assertTrue(net.training)
          net.eval()
          self.assertFalse(net.training)
          net.train()
          self.assertTrue(net.training)
          net.eval()
  
          # Try calling the additional methods we registered.
          biased_input = torch.randn(4, 5)
          output_before = net.forward(biased_input)
          bias = net.get_bias().clone()
          self.assertEqual(list(bias.shape), [2])
          net.set_bias(bias + 1)
          self.assertEqual(net.get_bias(), bias + 1)
          output_after = net.forward(biased_input)
  
          self.assertNotEqual(output_before, output_after)
  
          # Try accessing parameters
          self.assertEqual(len(net.parameters()), 2)
          np = net.named_parameters()
          self.assertEqual(len(np), 2)
          self.assertIn("fc.weight", np)
          self.assertIn("fc.bias", np)
  
          self.assertEqual(len(net.buffers()), 1)
          nb = net.named_buffers()
          self.assertEqual(len(nb), 1)
          self.assertIn("buf", nb)
          self.assertEqual(nb[0][1], torch.eye(5))
  
      def test_cpp_frontend_module_has_up_to_date_attributes(self):
          extension = torch.utils.cpp_extension.load(
              name="cpp_frontend_extension",
              sources="cpp_extensions/cpp_frontend_extension.cpp",
              verbose=True,
          )
  
          net = extension.Net(5, 2)
  
          self.assertEqual(len(net._parameters), 0)
          net.add_new_parameter("foo", torch.eye(5))
          self.assertEqual(len(net._parameters), 1)
  
          self.assertEqual(len(net._buffers), 1)
          net.add_new_buffer("bar", torch.eye(5))
          self.assertEqual(len(net._buffers), 2)
  
          self.assertEqual(len(net._modules), 1)
          net.add_new_submodule("fc2")
          self.assertEqual(len(net._modules), 2)
  
      @unittest.skipIf(not (TEST_CUDA or TEST_ROCM), "CUDA not found")
      def test_cpp_frontend_module_python_inter_op_with_cuda(self):
          extension = torch.utils.cpp_extension.load(
              name="cpp_frontend_extension",
              sources="cpp_extensions/cpp_frontend_extension.cpp",
              verbose=True,
          )
  
          net = extension.Net(5, 2)
          for p in net.parameters():
              self.assertTrue(p.device.type == "cpu")
          cpu_parameters = [p.clone() for p in net.parameters()]
  
          device = torch.device("cuda", 0)
          net.to(device)
  
          for i, p in enumerate(net.parameters()):
              self.assertTrue(p.device.type == "cuda")
              self.assertTrue(p.device.index == 0)
              self.assertEqual(cpu_parameters[i], p)
  
          net.cpu()
          net.add_new_parameter("a", torch.eye(5))
          net.add_new_parameter("b", torch.eye(5))
          net.add_new_buffer("c", torch.eye(5))
          net.add_new_buffer("d", torch.eye(5))
          net.add_new_submodule("fc2")
          net.add_new_submodule("fc3")
  
          for p in net.parameters():
              self.assertTrue(p.device.type == "cpu")
  
          net.cuda()
  
          for p in net.parameters():
              self.assertTrue(p.device.type == "cuda")
  
      def test_returns_shared_library_path_when_is_python_module_is_true(self):
          source = """
          #include <torch/script.h>
          torch::Tensor func(torch::Tensor x) { return x; }
          static torch::RegisterOperators r("test::func", &func);
          """
          torch.utils.cpp_extension.load_inline(
              name="is_python_module",
              cpp_sources=source,
              functions="func",
              verbose=True,
              is_python_module=False,
          )
          self.assertEqual(torch.ops.test.func(torch.eye(5)), torch.eye(5))
  
      def test_set_default_type_also_changes_aten_default_type(self):
          module = torch.utils.cpp_extension.load_inline(
              name="test_set_default_type",
              cpp_sources="torch::Tensor get() { return torch::empty({}); }",
              functions="get",
              verbose=True,
          )
  
          initial_default = torch.get_default_dtype()
          try:
              self.assertEqual(module.get().dtype, initial_default)
              torch.set_default_dtype(torch.float64)
              self.assertEqual(module.get().dtype, torch.float64)
              torch.set_default_dtype(torch.float32)
              self.assertEqual(module.get().dtype, torch.float32)
              torch.set_default_dtype(torch.float16)
              self.assertEqual(module.get().dtype, torch.float16)
          finally:
              torch.set_default_dtype(initial_default)
  
      def test_compilation_error_formatting(self):
          # Test that the missing-semicolon error message has linebreaks in it.
          # This'll fail if the message has been munged into a single line.
          # It's hard to write anything more specific as every compiler has it's own
          # error formatting.
          with self.assertRaises(RuntimeError) as e:
              torch.utils.cpp_extension.load_inline(
                  name="test_compilation_error_formatting",
                  cpp_sources="int main() { return 0 }")
          pattern = r'.*(\\n|\\r).*'
          self.assertNotRegex(str(e), pattern)
  
      def test_warning(self):
          # Note: the module created from this source will include the py::key_error
          # symbol. But because of visibility and the fact that it lives in a
          # different compilation unit than pybind, this trips up ubsan even though
          # it is fine. "ubsan.supp" thus needs to contain "vptr:warn_mod.so".
          source = '''
          // error_type:
          // 0: no error
          // 1: torch::TypeError
          // 2: python_error()
          // 3: py::error_already_set
          at::Tensor foo(at::Tensor x, int error_type) {
              std::ostringstream err_stream;
              err_stream << "Error with "  << x.type();
  
              TORCH_WARN(err_stream.str());
              if(error_type == 1) {
                  throw torch::TypeError(err_stream.str().c_str());
              }
              if(error_type == 2) {
                  PyObject* obj = PyTuple_New(-1);
                  TORCH_CHECK(!obj);
                  // Pretend it was caught in a different thread and restored here
                  auto e = python_error();
                  e.persist();
                  e.restore();
                  throw e;
              }
              if(error_type == 3) {
                  throw py::key_error(err_stream.str());
              }
              return x.cos();
          }
          '''
  
          # Ensure double type for hard-coded c name below
          t = torch.rand(2).double()
          cpp_tensor_name = r"CPUDoubleType"
  
          # Without error handling, the warnings cannot be catched
          warn_mod = torch.utils.cpp_extension.load_inline(name='warn_mod',
                                                           cpp_sources=[source],
                                                           functions=['foo'],
                                                           with_pytorch_error_handling=False)
  
          with warnings.catch_warnings(record=True) as w:
              warn_mod.foo(t, 0)
              self.assertEqual(len(w), 0)
  
              with self.assertRaisesRegex(TypeError, t.type()):
                  warn_mod.foo(t, 1)
              self.assertEqual(len(w), 0)
  
              with self.assertRaisesRegex(SystemError, "bad argument to internal function"):
                  warn_mod.foo(t, 2)
              self.assertEqual(len(w), 0)
  
              with self.assertRaisesRegex(KeyError, cpp_tensor_name):
                  warn_mod.foo(t, 3)
              self.assertEqual(len(w), 0)
  
  
          warn_mod = torch.utils.cpp_extension.load_inline(name='warn_mod',
                                                           cpp_sources=[source],
                                                           functions=['foo'],
                                                           with_pytorch_error_handling=True)
  
  
          with warnings.catch_warnings(record=True) as w:
              # Catched with no error should be detected
              warn_mod.foo(t, 0)
              self.assertEqual(len(w), 1)
  
              # Catched with cpp error should also be detected
              with self.assertRaisesRegex(TypeError, t.type()):
                  warn_mod.foo(t, 1)
              self.assertEqual(len(w), 2)
  
              # Catched with python error should also be detected
              with self.assertRaisesRegex(SystemError, "bad argument to internal function"):
                  warn_mod.foo(t, 2)
              self.assertEqual(len(w), 3)
  
              # Catched with pybind error should also be detected
              # Note that there is no type name translation for pybind errors
              with self.assertRaisesRegex(KeyError, cpp_tensor_name):
                  warn_mod.foo(t, 3)
              self.assertEqual(len(w), 4)
  
          # Make sure raising warnings are handled properly
          with warnings.catch_warnings(record=True) as w:
              warnings.simplefilter("error")
  
              # No error, the warning should raise
              with self.assertRaisesRegex(UserWarning, t.type()):
                  warn_mod.foo(t, 0)
              self.assertEqual(len(w), 0)
  
              # Another error happened, the warning is ignored
              with self.assertRaisesRegex(TypeError, t.type()):
                  warn_mod.foo(t, 1)
              self.assertEqual(len(w), 0)
  
      def test_autograd_from_cpp(self):
          source = '''
          void run_back(at::Tensor x) {
              x.backward({});
          }
  
          void run_back_no_gil(at::Tensor x) {
              pybind11::gil_scoped_release no_gil;
              x.backward({});
          }
          '''
  
          class MyFn(torch.autograd.Function):
              @staticmethod
              def forward(ctx, x):
                  return x.clone()
  
              @staticmethod
              def backward(ctx, gx):
                  return gx
  
          test_backward_deadlock = torch.utils.cpp_extension.load_inline(name='test_backward_deadlock',
                                                                         cpp_sources=[source],
                                                                         functions=['run_back', 'run_back_no_gil'],)
  
          # This used to deadlock
          inp = torch.rand(20, requires_grad=True)
          loss = MyFn.apply(inp).sum()
          with self.assertRaisesRegex(RuntimeError, "The autograd engine was called while holding the GIL."):
              test_backward_deadlock.run_back(loss)
  
          inp = torch.rand(20, requires_grad=True)
          loss = MyFn.apply(inp).sum()
          test_backward_deadlock.run_back_no_gil(loss)
  
      def test_custom_compound_op_autograd(self):
          # Test that a custom compound op (i.e. a custom op that just calls other aten ops)
          # correctly returns gradients of those other ops
  
          source = """
          #include <torch/library.h>
          torch::Tensor my_add(torch::Tensor x, torch::Tensor y) {
            return x + y;
          }
          TORCH_LIBRARY(my, m) {
              m.def("add", &my_add);
          }
          """
  
          torch.utils.cpp_extension.load_inline(
              name="is_python_module",
              cpp_sources=source,
              verbose=True,
              is_python_module=False,
          )
  
          a = torch.randn(5, 5, requires_grad=True)
          b = torch.randn(5, 5, requires_grad=True)
  
          gradcheck(torch.ops.my.add, [a, b], eps=1e-2)
  
      @staticmethod
      def _crash_handler_test_process(stderr_file, destination):
          # Code to enable dumps and trigger a segfault
          if sys.platform == "win32":
              destination = destination.replace("\\", "\\\\")
              csrc = textwrap.dedent(f"""
              #include <torch/torch.h>
              #include <locale>
              #include <iostream>
              #include <codecvt>
              #include <string>
  
              int fail() {{
                  std::wstring_convert<std::codecvt_utf8_utf16<wchar_t>> converter;
                  std::string narrow("{destination}");
                  std::wstring wide = converter.from_bytes(narrow);
                  torch::crash_handler::enable_minidumps(wide.c_str());
  
                  volatile int* bad = nullptr;
                  return *bad;
              }}
              """)
          else:
              csrc = textwrap.dedent(f"""
              #include <torch/torch.h>
  
              int fail() {{
                  torch::crash_handler::enable_minidumps("{destination}");
  
                  volatile int* bad = nullptr;
                  return *bad;
              }}
              """)
  
          # Some special stuff to overwrite stderr for a C++ extension
          # Copied from: https://stackoverflow.com/questions/8804893/redirect-stdout-from-python-for-c-calls
          sys.stdout.flush()
          newstdout = os.dup(2)
          devnull = os.open(stderr_file, os.O_WRONLY)
          os.dup2(devnull, 2)
          os.close(devnull)
          sys.stdout = os.fdopen(newstdout, 'w')
  
          module = torch.utils.cpp_extension.load_inline(
              name="segfault",
              cpp_sources=csrc,
              functions=["fail"],
          )
          module.fail()
  
      @unittest.skipIf(TEST_WITH_ASAN, "ASAN disables the crash handler's signal handler")
      @unittest.skipIf(not has_breakpad(), "Built without breakpad")
      @unittest.skipIf(os.environ.get("TEST_CONFIG") == "force_on_cpu", "fails on force_on_cpu config, tracked w/ #65253")
      def test_crash_handler(self):
          with tempfile.TemporaryDirectory() as temp_dir, tempfile.NamedTemporaryFile(delete=not sys.platform == "win32") as stderr:
              # Use multiprocessing to spin up a separate process to make catching
              # the segfault easier
              p = Process(target=self._crash_handler_test_process, args=(stderr.name, temp_dir))
              p.start()
              p.join()
  
              with open(stderr.name) as f:
                  result = f.read().strip()
  
              # Check that the signal handler was called
              self.assertTrue(result.startswith(f"Wrote minidump to {temp_dir}"))
  
              with open(result.replace("Wrote minidump to ", ""), "rb") as dump_file:
                  dump_bytes = dump_file.read()
  
                  # Check that the file has the correct magic number
                  self.assertEqual(b"MDMP", dump_bytes[0:4])
  
  
  
  if __name__ == "__main__":
      common.run_tests()
  