# mypy: allow-untyped-defs # Owner(s): ["module: unknown"] import os import random import re import shutil import subprocess import sys import tempfile import textwrap import traceback import unittest import warnings from typing import Any, Dict, List import torch import torch.cuda import torch.nn as nn import torch.utils.cpp_extension import torch.utils.data from torch.autograd._functions.utils import check_onnx_broadcast from torch.onnx.symbolic_opset9 import _prepare_onnx_paddings from torch.testing._internal.common_cuda import TEST_MULTIGPU from torch.testing._internal.common_device_type import ( instantiate_device_type_tests, onlyCPU, ops, ) from torch.testing._internal.common_methods_invocations import op_db from torch.testing._internal.common_utils import ( # type: ignore[attr-defined] IS_FBCODE, IS_SANDCASTLE, IS_WINDOWS, load_tests, ) from torch.utils._device import set_device from torch.utils._pytree import tree_all_only, tree_any from torch.utils._traceback import ( CapturedTraceback, format_traceback_short, report_compile_source_on_error, ) from torch.utils.checkpoint import ( _infer_device_type, checkpoint, checkpoint_sequential, get_device_states, ) from torch.utils.data import DataLoader # load_tests from torch.testing._internal.common_utils is used to automatically filter tests for # sharding on sandcastle. This line silences flake warnings load_tests = load_tests HAS_CUDA = torch.cuda.is_available() from torch.testing._internal.common_utils import run_tests, TestCase class RandomDatasetMock(torch.utils.data.Dataset): def __getitem__(self, index): return torch.tensor([torch.rand(1).item(), random.uniform(0, 1)]) def __len__(self): return 1000 class TestCheckpoint(TestCase): # This runs checkpoint_sequential on each of the nets in # module_lists_to_compare, and compares them against the uncheckpointed model. # To compare, it checks outputs as well as input gradients and parameter gradients def _check_checkpoint_sequential( self, model, module_lists_to_compare, num_chunks, input, use_reentrant, ): # not checkpointed out = model(input) out_not_checkpointed = out.detach().clone() model.zero_grad() out.sum().backward() grad_not_checkpointed = { name: param.grad.detach().clone() for name, param in model.named_parameters() } input_grad_not_checkpointed = input.grad.detach().clone() for model_to_compare in module_lists_to_compare: # checkpointed model by passing list of modules detached = input.detach() detached.requires_grad = True # pass list of modules to checkpoint out = checkpoint_sequential( model_to_compare, num_chunks, detached, use_reentrant=use_reentrant ) out_checkpointed = out.detach().clone() model.zero_grad() out.sum().backward() grad_checkpointed = { name: param.grad.detach().clone() for name, param in model.named_parameters() } input_grad_checkpointed = detached.grad.detach().clone() # compare outputs as well as the gradients of input and parameters self.assertEqual(out_checkpointed, out_not_checkpointed) self.assertEqual(input_grad_not_checkpointed, input_grad_checkpointed) for name in grad_checkpointed: self.assertEqual(grad_checkpointed[name], grad_not_checkpointed[name]) # Test whether checkpoint is being triggered or not. For this, we check # the number of times forward pass happens def test_checkpoint_trigger(self): class Net(nn.Module): def __init__(self) -> None: super().__init__() self.counter = 0 def forward(self, input_var): self.counter += 1 # For reentrant, need to have autograd actually # pack a tensor to trigger recomp ret = input_var * torch.tensor(2.0) return ret # checkpointed for use_reentrant in [True, False]: with self.subTest(use_reentrant=use_reentrant): modules = [Net() for _ in range(10)] for m in modules: self.assertEqual(m.counter, 0) input_var = torch.randn(3, 4, requires_grad=True) out = checkpoint_sequential( modules, 2, input_var, use_reentrant=use_reentrant ) for m in modules: self.assertEqual(m.counter, 1) out.sum().backward() for m in modules[: (len(modules) // 2)]: self.assertEqual(m.counter, 2) for m in modules[(len(modules) // 2) :]: self.assertEqual(m.counter, 1) def test_checkpoint_valid(self): model = nn.Sequential( nn.Linear(100, 50), nn.ReLU(), nn.Linear(50, 20), nn.ReLU(), nn.Linear(20, 5), nn.ReLU(), ) input_var = torch.randn(1, 100, requires_grad=True) # checkpointed chunks = 2 modules = list(model.children()) out = checkpoint_sequential(modules, chunks, input_var, use_reentrant=True) with self.assertRaisesRegex( RuntimeError, "torch.utils.checkpoint is incompatible" ): torch.autograd.grad( outputs=[out], grad_outputs=[torch.ones(1, 5)], inputs=[input_var], create_graph=True, ) # works with use_reentrant=False, and grads are the same out = model(input_var) grads_no_checkpoint = torch.autograd.grad( outputs=[out], grad_outputs=[torch.ones(1, 5)], inputs=[input_var], create_graph=True, ) out_checkpoint = checkpoint_sequential( modules, chunks, input_var, use_reentrant=False ) # check outputs are the same self.assertEqual(out_checkpoint, out) grads_checkpoint = torch.autograd.grad( outputs=[out_checkpoint], grad_outputs=[torch.ones(1, 5)], inputs=[input_var], create_graph=True, ) self.assertEqual(grads_no_checkpoint, grads_checkpoint) def test_checkpoint(self): for use_reentrant in [True, False]: with self.subTest(use_reentrant=use_reentrant): model = nn.Sequential( nn.Linear(100, 50), nn.ReLU(), nn.Linear(50, 20), nn.ReLU(), nn.Linear(20, 5), nn.ReLU(), ) # Compare uncheckpointed model with its checkpointed counterparts # In addition to running checkpoint_sequential on the nn.Sequential # instance, we also run the function on the list of functions within # the module. self._check_checkpoint_sequential( model, [list(model.children()), model], 2, torch.randn(1, 100, requires_grad=True), use_reentrant=use_reentrant, ) def test_checkpoint_module_list(self): class ModuleListNet(nn.Module): def __init__(self) -> None: super().__init__() module_list = [ nn.Linear(100, 50), nn.ReLU(), nn.Linear(50, 20), nn.ReLU(), nn.Linear(20, 5), nn.ReLU(), ] self.module_list = nn.ModuleList(module_list) def forward(self, input): for layer in self.module_list: input = layer(input) return input for use_reentrant in [True, False]: with self.subTest(use_reentrant=use_reentrant): model = ModuleListNet() # Compare uncheckpointed model with its checkpointed counterparts. self._check_checkpoint_sequential( model, [list(model.module_list.children()), model.module_list], 2, torch.randn(1, 100, requires_grad=True), use_reentrant=use_reentrant, ) def test_checkpoint_sequential_deprecated_multiple_args(self): class Two(nn.Module): def forward(self, a, b): return a, b model = nn.Sequential(Two()) a = torch.randn(1, 100, requires_grad=True) b = torch.randn(1, 100, requires_grad=True) for use_reentrant in [True, False]: with self.subTest(use_reentrant=use_reentrant): with self.assertRaises(TypeError): checkpoint_sequential(model, 1, a, b) # type: ignore[call-arg] def test_checkpoint_sequential_deprecated_no_args(self): class Noop(nn.Module): def forward(self): pass model = nn.Sequential(Noop()) for use_reentrant in [True, False]: with self.subTest(use_reentrant=use_reentrant): with self.assertRaises(TypeError): checkpoint_sequential(model, 1) # type: ignore[call-arg] def test_checkpoint_rng_cpu(self): for _ in range(5): inp = torch.randn(20000, device="cpu").requires_grad_() phase1 = torch.nn.Dropout() phase2 = torch.nn.Dropout() def run_fn(input): return phase2(input) state = torch.get_rng_state() out = phase1(inp) out = checkpoint(run_fn, out, use_reentrant=True) out.sum().backward() grad_with_checkpointing = inp.grad torch.set_rng_state(state) inp.grad = None out = phase1(inp) out = run_fn(out) out.sum().backward() grad_no_checkpointing = inp.grad self.assertEqual(grad_with_checkpointing, grad_no_checkpointing) @unittest.skipIf(not HAS_CUDA, "No CUDA") def test_checkpoint_rng_cuda(self): for _ in range(5): inp = torch.randn(20000, device="cuda").requires_grad_() phase1 = torch.nn.Dropout() phase2 = torch.nn.Dropout() def run_fn(input): return phase2(input) state = torch.cuda.get_rng_state() out = phase1(inp) out = checkpoint(run_fn, out, use_reentrant=True) out.sum().backward() grad_with_checkpointing = inp.grad torch.cuda.set_rng_state(state) inp.grad = None out = phase1(inp) out = run_fn(out) out.sum().backward() grad_no_checkpointing = inp.grad self.assertEqual(grad_with_checkpointing, grad_no_checkpointing) @unittest.skipIf(not HAS_CUDA, "No CUDA") def test_checkpoint_not_preserve_rng_state_and_without_reentrant(self): inp = torch.randn(2, device="cuda").requires_grad_() layer = torch.nn.Dropout() def run_fn(input): return layer(input) out = checkpoint(run_fn, inp, use_reentrant=False, preserve_rng_state=False) out.sum().backward() # This should run without error def test_checkpoint_non_tensor(self): def run_fn(tensor1, tensor2): if tensor2 is None: return tensor1 return tensor1 + tensor2 input_var = torch.randn(1, 100, requires_grad=True) out = checkpoint(run_fn, input_var, None, use_reentrant=True) out.sum().backward() def test_checkpoint_non_tensor_inputs_outputs(self): def foo(t1, t2, scale, t3): t4 = t1 + t2 * t3 t5 = t1 * t2 + t3 t4 *= scale t5 *= scale return scale, t4, None, True, t5, "bar", t1 t1 = torch.rand(10, requires_grad=True) t2 = torch.rand(10, requires_grad=True) t3 = torch.rand(10) scale = random.randint(0, 10) res = checkpoint(foo, t1, t2, scale, t3, use_reentrant=True) self.assertEqual(scale, res[0]) self.assertEqual((t1 + t2 * t3) * scale, res[1]) self.assertEqual(None, res[2]) self.assertEqual(True, res[3]) self.assertEqual((t1 * t2 + t3) * scale, res[4]) self.assertEqual("bar", res[5]) self.assertEqual(t1, res[6]) # Validate running backward. res[1].sum().backward(retain_graph=True) res[4].sum().backward(retain_graph=True) res[6].sum().backward() with self.assertRaisesRegex( RuntimeError, "Trying to backward through the graph a second time" ): res[6].sum().backward() t1_grad = t1.grad t2_grad = t2.grad # Reset grads, run without checkpoint and validate we receive same grads. t1.grad = None t2.grad = None res = foo(t1, t2, scale, t3) torch.autograd.backward([res[1].sum(), res[4].sum(), res[6].sum()]) self.assertEqual(t1.grad, t1_grad) self.assertEqual(t2.grad, t2_grad) def test_checkpoint_no_tensors(self): def foo(t1, t2, scale, t3): t4 = t1 + t2 * t3 t5 = t1 * t2 + t3 t4 *= scale t5 *= scale return scale, t4, None, True, t5, "bar", t1 t1 = random.random() t2 = random.random() t3 = random.random() scale = random.randint(0, 10) res = checkpoint(foo, t1, t2, scale, t3, use_reentrant=True) self.assertEqual(scale, res[0]) self.assertEqual((t1 + t2 * t3) * scale, res[1]) self.assertEqual(None, res[2]) self.assertEqual(True, res[3]) self.assertEqual((t1 * t2 + t3) * scale, res[4]) self.assertEqual("bar", res[5]) self.assertEqual(t1, res[6]) def test_checkpoint_partial_grad(self): def run_fn(tensor1, tensor2): # tensor 2 is used for other application logic return tensor1, tensor2 input_var = torch.randn(1, 4, requires_grad=True) input_var2 = torch.randn(1, 4, requires_grad=False) out = checkpoint(run_fn, input_var, input_var2, use_reentrant=True) out[0].sum().backward() def run_fn2(tensor1, tensor2): return tensor1 input_var = torch.randn(1, 4, requires_grad=False) input_var2 = torch.randn(1, 4, requires_grad=True) with self.assertRaisesRegex( RuntimeError, r"none of output has requires_grad=True, this checkpoint\(\) is not necessary", ): out = checkpoint(run_fn2, input_var, input_var2, use_reentrant=True) out.sum().backward() @unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") def test_checkpointing_without_reentrant_early_free(self): # I don't know how to check if the temporary saved variable buffer # get de-allocated directly. So using cuda memory usage as a proxy def _do_test(fn, should_free): stats: List[int] = [] def track(x, idx): # Track that at each step of the backward, some Tensor were # de-allocated (which correspond to the checkpoint storage being # emptied at each step) def hook(_unused): self.assertEqual(len(stats), idx) torch.cuda.synchronize() stats.append(torch.cuda.memory_allocated()) if idx > 0: if should_free: self.assertLess(stats[idx], stats[idx - 1]) else: self.assertEqual(stats[idx], stats[idx - 1]) x.register_hook(hook) def test_fn(x): # The main property of this function is that it contains multiple # operations that save gradients in a chain. x = x**2 track(x, 2) x = x**2 track(x, 1) x = x**2 track(x, 0) x = x**2 return x.sum() fn(test_fn) return stats x = torch.zeros(10, device="cuda", requires_grad=True) x.grad = torch.zeros_like(x) # In a regular backward, buffers get eagerly freed non_retain_stats = _do_test(lambda fn: fn(x).backward(), True) # In a retain_grad backward, buffers get preserved _unused_retain_stats = _do_test( lambda fn: fn(x).backward(retain_graph=True), False ) # In a regular backward with checkpoint, buffers get eagerly freed checkpoint_non_retain_stats = _do_test( lambda fn: checkpoint(fn, x, use_reentrant=False).backward(), True ) # In a retain_grad backward with checkpoint, buffers get eagerly freed checkpoint_retain_stats = _do_test( lambda fn: checkpoint(fn, x, use_reentrant=False).backward( retain_graph=True ), True, ) self.assertEqual(non_retain_stats, checkpoint_non_retain_stats) self.assertEqual(non_retain_stats, checkpoint_retain_stats) @unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported") def test_get_device_states_recursive(self): inp = { "foo": torch.rand(10, device="cuda:0"), "bar": [torch.rand(10, device="cuda:1")], } device_ids, device_states = get_device_states(inp) self.assertEqual(2, len(device_ids)) self.assertEqual(2, len(device_states)) self.assertEqual(0, device_ids[0]) self.assertEqual(1, device_ids[1]) self.assertTrue(isinstance(device_states[0], torch.Tensor)) self.assertTrue(isinstance(device_states[1], torch.Tensor)) def test_infer_device_state_recursive_meta(self): inp = {"foo": torch.rand(10, device="meta")} device_type = _infer_device_type(inp) self.assertEqual("meta", device_type) @unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported") def test_infer_device_state_recursive_multi_cuda(self): # Check that no warning is issued for either cuda:0, cuda:1 or # cuda:0, cuda:0 cases since they are both the same device type inp = { "foo": torch.rand(10, device="cuda:0"), "bar": [torch.rand(10, device="cuda:1")], } with warnings.catch_warnings(): warnings.simplefilter("error") device_type = _infer_device_type(inp) self.assertEqual("cuda", device_type) inp = { "foo": torch.rand(10, device="cuda:0"), "bar": [torch.rand(10, device="cuda:0")], } with warnings.catch_warnings(): warnings.simplefilter("error") device_type = _infer_device_type(inp) self.assertEqual("cuda", device_type) # Check that a warning is issued for cuda:0, meta and that it includes # device type information inp = { "foo": torch.rand(10, device="cuda:0"), "bar": [torch.rand(10, device="meta")], } with warnings.catch_warnings(record=True) as w: device_type = _infer_device_type(inp) self.assertEqual("cuda", device_type) self.assertEqual(len(w), 1) warning_msg = str(w[-1].message) self.assertTrue( "Tensor arguments, excluding CPU tensors, are detected on at least two types of devices" in warning_msg ) self.assertTrue("Device types: ['cuda', 'meta']" in warning_msg) self.assertTrue("first device type: cuda" in warning_msg) class TestDataLoaderUtils(TestCase): MAX_TIMEOUT_IN_SECOND = 300 def setUp(self): super().setUp() self.dataset = torch.randn(5, 3, 3, 2) self.batch_size = 3 def test_random_seed(self): def run(): dataloader = torch.utils.data.DataLoader( RandomDatasetMock(), batch_size=2, num_workers=4, shuffle=True, timeout=self.MAX_TIMEOUT_IN_SECOND, ) return next(iter(dataloader)) torch.manual_seed(2018) x1 = run() torch.manual_seed(2018) x2 = run() self.assertEqual(x1, x2) def test_single_keep(self): # self.dataset is a Tensor here; technically not a valid input because # not a Dataset subclass, but needs to stay working so add ignore's # for type checking with mypy dataloader: DataLoader = DataLoader( self.dataset, # type: ignore[arg-type] batch_size=self.batch_size, num_workers=0, drop_last=False, ) dataiter = iter(dataloader) self.assertEqual(len(list(dataiter)), 2) def test_single_drop(self): dataloader: DataLoader = DataLoader( self.dataset, # type: ignore[arg-type] batch_size=self.batch_size, num_workers=0, drop_last=True, ) dataiter = iter(dataloader) self.assertEqual(len(list(dataiter)), 1) @unittest.skip( "FIXME: Intermittent CUDA out-of-memory error on Windows and time-out under ASAN" ) def test_multi_keep(self): dataloader: DataLoader = DataLoader( self.dataset, # type: ignore[arg-type] batch_size=self.batch_size, num_workers=2, drop_last=False, timeout=self.MAX_TIMEOUT_IN_SECOND, ) dataiter = iter(dataloader) self.assertEqual(len(list(dataiter)), 2) def test_multi_drop(self): dataloader: DataLoader = DataLoader( self.dataset, # type: ignore[arg-type] batch_size=self.batch_size, num_workers=2, drop_last=True, timeout=self.MAX_TIMEOUT_IN_SECOND, ) dataiter = iter(dataloader) self.assertEqual(len(list(dataiter)), 1) test_dir = os.path.abspath(os.path.dirname(str(__file__))) @unittest.skipIf( "SKIP_TEST_BOTTLENECK" in os.environ.keys(), "SKIP_TEST_BOTTLENECK is set" ) class TestBottleneck(TestCase): def _run(self, command, timeout=30): """Returns (return-code, stdout, stderr)""" import subprocess p = subprocess.Popen( command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, ) try: output, err = p.communicate(timeout=timeout) except subprocess.TimeoutExpired: p.kill() output, err = p.communicate() rc = p.returncode output_str = output.decode("ascii") err_str = err.decode("ascii") return (rc, output_str, err_str) def _run_bottleneck(self, test_file, scriptargs=""): curdir = os.path.dirname(os.path.abspath(__file__)) filepath = f"{curdir}/{test_file}" if scriptargs != "": scriptargs = f" {scriptargs}" rc, out, err = self._run( f"{sys.executable} -m torch.utils.bottleneck {filepath}{scriptargs}" ) return rc, out, err def _check_run_args(self): # Check that this fails due to missing args rc, out, err = self._run_bottleneck("bottleneck_test/test_args.py") self.assertEqual( rc, 2, atol=0, rtol=0, msg=self._fail_msg("Missing args should error", out + err), ) # This should succeed rc, out, err = self._run_bottleneck( "bottleneck_test/test_args.py", "--foo foo --bar bar" ) self.assertEqual( rc, 0, atol=0, rtol=0, msg=self._fail_msg("Should pass args to script", out + err), ) def _fail_msg(self, msg, output): return f"{msg}, output was:\n{output}" def _check_environment_summary(self, output): results = re.search("Environment Summary", output) self.assertIsNotNone( results, self._fail_msg("Should have Environment Summary", output) ) # Up to five lines away from the heading, there should be the version number results = re.search( r"Environment Summary.*(\n.*){,5}\nPyTorch \d+\.\d+", output ) self.assertIsNotNone( results, self._fail_msg("Should have PyTorch version", output) ) def _check_cprof_summary(self, output): results = re.search("cProfile output", output) self.assertIsNotNone( results, self._fail_msg("Should have cProfile output", output) ) # This assumes that after the cProfile output section we have # the autograd profiler output results = re.search( r"cProfile output.*(\n.*){6,50}\n.*autograd profiler output", output ) self.assertIsNotNone( results, self._fail_msg( "Distance between cProfile and autograd prof out not in [6, 50] lines", output, ), ) def _check_autograd_summary(self, output): results = re.search("autograd profiler output", output) self.assertIsNotNone( results, self._fail_msg("Should have autograd profiler output", output) ) # This assumes that after the autograd profiler output is the end of the # output. results = re.search(r"autograd profiler output.*(\n.*){6,100}", output) self.assertIsNotNone( results, self._fail_msg( "Distance between autograd prof output and end of output not in [6, 100] lines", output, ), ) def _check_cuda(self, output): if HAS_CUDA: results = re.search("CUDA mode", output) self.assertIsNotNone( results, self._fail_msg("Should tell users CUDA", output) ) else: results = re.search("CUDA mode", output) self.assertIsNone( results, self._fail_msg("Should not tell users about CUDA", output) ) @unittest.skipIf(HAS_CUDA, "CPU-only test") def test_bottleneck_cpu_only(self): rc, out, err = self._run_bottleneck("bottleneck_test/test.py") self.assertEqual(rc, 0, msg=f"Run failed with\n{err}") self._check_run_args() self._check_environment_summary(out) self._check_autograd_summary(out) self._check_cprof_summary(out) self._check_cuda(out) @unittest.skipIf(not HAS_CUDA, "No CUDA") def test_bottleneck_cuda(self): rc, out, err = self._run_bottleneck("bottleneck_test/test_cuda.py") self.assertEqual(rc, 0, msg=f"Run failed with\n{err}") self._check_run_args() self._check_environment_summary(out) self._check_autograd_summary(out) self._check_cprof_summary(out) self._check_cuda(out) from torch.utils.collect_env import get_pretty_env_info @unittest.skipIf(IS_FBCODE, "runs pip which is not available internally") class TestCollectEnv(TestCase): def test_smoke(self): info_output = get_pretty_env_info() self.assertTrue(info_output.count("\n") >= 17) class TestONNXUtils(TestCase): def test_prepare_onnx_paddings(self): sizes = [2, 3, 4] pad = [1, 2, 3, 4] paddings = _prepare_onnx_paddings(len(sizes), pad) self.assertEqual(paddings, [0, 3, 1, 0, 4, 2]) def test_check_onnx_broadcast(self): def try_check_onnx_broadcast(dims1, dims2, expect_broadcast, expect_fail): broadcast = True fail = False try: broadcast = check_onnx_broadcast(dims1, dims2) except ValueError: fail = True self.assertEqual(broadcast, expect_broadcast) self.assertEqual(fail, expect_fail) # Case 1, check the case when len(dims1) < len(dims2) and numel(dims2) > 1 dims1 = [3, 4] dims2 = [2, 3, 4] try_check_onnx_broadcast(dims1, dims2, True, True) # Case 2, check the case when len(dims1) < len(dims2) and numel(dims2) == 1 dims1 = [3, 4] dims2 = [1, 1, 1] try_check_onnx_broadcast(dims1, dims2, True, False) # Case 3, check the case when len(dims1) > len(dims2) and numel(dims2) == 1 dims1 = [1, 1] dims2 = [1] try_check_onnx_broadcast(dims1, dims2, True, False) # Case 4, check the case when len(dims1) > len(dims2) and dims1[x:] == dims2 dims1 = [2, 3, 4] dims2 = [3, 4] try_check_onnx_broadcast(dims1, dims2, True, False) # Case 5, check the case when len(dims1) > len(dims2), but dims1[x:] != dims2 dims1 = [2, 3, 4] dims2 = [1, 4] try_check_onnx_broadcast(dims1, dims2, True, True) # Case 6, check the equal case, no broadcast dims1 = [3, 4] dims2 = [3, 4] try_check_onnx_broadcast(dims1, dims2, False, False) # Case 7, check the case when len(dims1) == len(dims2), but dims1 != dims2 dims1 = [3, 4] dims2 = [1, 4] try_check_onnx_broadcast(dims1, dims2, True, True) # Case 8, check the case when len(dims1) == len(dims2) and numel(s2) == 1 dims1 = [3, 4] dims2 = [1, 1] try_check_onnx_broadcast(dims1, dims2, True, False) class TestHipify(TestCase): def test_import_hipify(self): from torch.utils.hipify import hipify_python # noqa: F401 class TestHipifyTrie(TestCase): def setUp(self): self.trie = torch.utils.hipify.hipify_python.Trie() def test_add_and_search_trie(self): self.trie.add("banana") self.assertTrue(self.trie.search("banana")) self.assertFalse(self.trie.search("ban")) self.assertFalse(self.trie.search("dog")) def test_add_multiple_and_search_trie(self): words_to_add = ["banana", "apple", "orange"] for word in words_to_add: self.trie.add(word) for word in words_to_add: self.assertTrue(self.trie.search(word)) for word in ["ban", "dog", "okay", "app"]: self.assertFalse(self.trie.search(word)) def test_quote_escape(self): orig_chars = ["*", "[", ".", "+", "a", "z", "-"] quoted_strs = ["\\*", "\\[", "\\.", "\\+", "a", "z", "\\-"] for i in range(len(orig_chars)): self.assertEqual(self.trie.quote(orig_chars[i]), quoted_strs[i]) def test_export_trie_to_regex(self): words_to_add = [ "__CUDACC__", "CUDA_ERROR_CONTEXT_ALREADY_CURRENT", "CUDA_ERROR_ARRAY_IS_MAPPED", "CUDA_ERROR_NOT_MAPPED", "CUDA_ERROR_INVALID_SOURCE", ] for word in words_to_add: self.trie.add(word) regex = self.trie.export_to_regex() expected_regex = r"(?:CUDA_ERROR_(?:ARRAY_IS_MAPPED|CONTEXT_ALREADY_CURRENT|INVALID_SOURCE|NOT_MAPPED)|__CUDACC__)" self.assertEqual(regex, expected_regex) def test_prefix_words_export_trie_to_regex(self): # test case where some nodes have both children and are also leaf nodes. words_to_add = ["apple", "app", "ban", "banana"] for word in words_to_add: self.trie.add(word) regex = self.trie.export_to_regex() expected_regex = r"(?:app(?:le)?|ban(?:ana)?)" self.assertEqual(regex, expected_regex) def test_single_export_trie_to_regex(self): words_to_add = ["cudaErrorInvalidMemcpyDirection"] for word in words_to_add: self.trie.add(word) regex = self.trie.export_to_regex() expected_regex = "cudaErrorInvalidMemcpyDirection" self.assertEqual(regex, expected_regex) def test_char_export_trie_to_regex(self): self.trie.add("a") self.assertEqual(self.trie.export_to_regex(), "a") self.trie.add("b") self.assertEqual(self.trie.export_to_regex(), "[ab]") def test_special_char_export_trie_to_regex(self): self.trie.add(r"c*") self.assertEqual(self.trie.export_to_regex(), r"c\*") class TestAssert(TestCase): def test_assert_true(self): # verify assertions work as expected # bool argument torch._assert(True, "foo") with self.assertRaisesRegex(AssertionError, "bar"): torch._assert(False, "bar") # tensor argument torch._assert(torch.tensor([True], dtype=torch.bool), "foo") with self.assertRaisesRegex(AssertionError, "bar"): torch._assert(torch.tensor([False], dtype=torch.bool), "bar") def test_assert_scriptable(self): class M(torch.nn.Module): def forward(self, x): torch._assert(x.sum() > 0, "foo") return x m = M() # scriptable ms = torch.jit.script(m) # data can be passed without errors x = torch.randn(4, 4).fill_(1.0) ms(x) with self.assertRaisesRegex(torch.jit.Error, "foo"): ms(torch.tensor([False], dtype=torch.bool)) @unittest.skipIf(IS_SANDCASTLE, "cpp_extension is OSS only") class TestStandaloneCPPJIT(TestCase): def test_load_standalone(self): build_dir = tempfile.mkdtemp() try: src_path = os.path.join(build_dir, "main.cpp") src = textwrap.dedent( """\ #include #include int main() { auto x = torch::eye(3); std::cout << x << std::endl; } """ ) with open(src_path, "w") as f: f.write(src) exec_path = torch.utils.cpp_extension.load( "standalone_load_test", src_path, build_directory=build_dir, is_python_module=False, is_standalone=True, ) ext = ".exe" if IS_WINDOWS else "" self.assertEqual( exec_path, os.path.join(build_dir, f"standalone_load_test{ext}") ) for shell in [True, False]: r = subprocess.run( [exec_path], shell=shell, stdout=subprocess.PIPE, ) self.assertEqual(r.returncode, 0) self.assertEqual( # Windows prints "\r\n" for newlines. textwrap.dedent(r.stdout.decode("utf-8")).replace("\r\n", "\n"), textwrap.dedent( """\ 1 0 0 0 1 0 0 0 1 [ CPUFloatType{3,3} ] """ ), ) finally: shutil.rmtree(build_dir) class DummyPrivateUse1Module: @staticmethod def is_available(): return True @staticmethod def is_autocast_enabled(): return True @staticmethod def get_autocast_dtype(): return torch.float16 @staticmethod def set_autocast_enabled(enable): pass @staticmethod def set_autocast_dtype(dtype): pass @staticmethod def get_amp_supported_dtype(): return [torch.float16] class TestExtensionUtils(TestCase): def tearDown(self): # Clean up backend_name = torch._C._get_privateuse1_backend_name() if hasattr(torch, backend_name): delattr(torch, backend_name) if f"torch.{backend_name}" in sys.modules: del sys.modules[f"torch.{backend_name}"] def test_external_module_register(self): # Built-in module with self.assertRaisesRegex(RuntimeError, "The runtime module of"): torch._register_device_module("cuda", torch.cuda) # Wrong device type with self.assertRaisesRegex(RuntimeError, "Expected one of cpu"): torch._register_device_module("dummmy", DummyPrivateUse1Module) with self.assertRaises(AttributeError): torch.privateuseone.is_available() # type: ignore[attr-defined] torch._register_device_module("privateuseone", DummyPrivateUse1Module) torch.privateuseone.is_available() # type: ignore[attr-defined] # No supporting for override with self.assertRaisesRegex(RuntimeError, "The runtime module of"): torch._register_device_module("privateuseone", DummyPrivateUse1Module) def test_external_module_register_with_renamed_backend(self): torch.utils.rename_privateuse1_backend("foo") with self.assertRaisesRegex(RuntimeError, "has already been set"): torch.utils.rename_privateuse1_backend("dummmy") custom_backend_name = torch._C._get_privateuse1_backend_name() self.assertEqual(custom_backend_name, "foo") with self.assertRaises(AttributeError): torch.foo.is_available() # type: ignore[attr-defined] with self.assertRaisesRegex(AssertionError, "Tried to use AMP with the"): with torch.autocast(device_type=custom_backend_name): pass torch._register_device_module("foo", DummyPrivateUse1Module) torch.foo.is_available() # type: ignore[attr-defined] with torch.autocast(device_type=custom_backend_name): pass self.assertEqual(torch._utils._get_device_index("foo:1"), 1) self.assertEqual(torch._utils._get_device_index(torch.device("foo:2")), 2) class TestRenderUtils(TestCase): def test_basic(self): self.assertExpectedInline( torch._utils.render_call(torch.sum, [torch.randn(100)], {"dim": 0}), """torch.sum(tensor([...], size=(100,)), dim=0)""", ) self.assertExpectedInline( torch._utils.render_call(torch.sum, [torch.randn(100, 100)], {"dim": 0}), """torch.sum(tensor([...], size=(100, 100)), dim=0)""", ) class TestDeviceUtils(TestCase): def test_basic(self): with torch.device("meta") as dev: x = torch.empty(3, 3) self.assertEqual(x.device.type, "meta") self.assertEqual(dev, torch.device("meta")) def test_decorator(self): @set_device("meta") def f(): return torch.empty(3, 3) self.assertEqual(f().device.type, "meta") def test_decorator_generator(self): @set_device("meta") def f(): yield torch.empty(3, 3) yield torch.empty(3, 3) r1, r2 = list(f()) self.assertEqual(r1.device.type, "meta") self.assertEqual(r2.device.type, "meta") def test_nn_module(self): with torch.device("meta"): m = nn.Linear(40, 50) self.assertEqual(m.weight.device.type, "meta") def test_set_default_device(self): try: torch.set_default_device("meta") r = torch.empty(2, 2) finally: torch.set_default_device(None) self.assertEqual(r.device.type, "meta") def test_get_default_device(self): torch.set_default_device("meta") self.assertEqual(torch.get_default_device().type, "meta") torch.set_default_device(None) @unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported") def test_get_default_device_more(self): torch.set_default_device("cuda") self.assertEqual(torch.get_default_device(), torch.tensor([]).device) torch.set_default_device(None) torch.set_default_device("cuda") torch.cuda.set_device("cuda:1") self.assertEqual(torch.get_default_device(), torch.tensor([]).device) torch.set_default_device(None) torch.set_default_device("cuda:1") self.assertEqual(torch.get_default_device(), torch.tensor([]).device) torch.set_default_device(None) @onlyCPU @ops(op_db) def test_device_mode_ops(self, device, dtype, op): func = op.get_op() samples = op.sample_inputs(device, dtype, requires_grad=False) for sample in samples: # Only test samples which don't have Tensor inputs. However, # we don't test the factory property on OpInfo as it is very, # very incomplete if tree_any( lambda x: isinstance(x, torch.Tensor), (sample.input, sample.args, sample.kwargs), ): continue # Many OpInfos will explicitly pass in a device. DeviceContext # will respect device if it is explicitly specified. To test # DeviceContext, we have to remove the device kwarg in this case. # NB: Can't pass None to sample_inputs, the function can't # handle it. kwargs = sample.kwargs.copy() kwargs.pop("device", None) with torch.device("meta"): r = func(sample.input, *sample.args, **kwargs) def is_meta_device(x: torch.Tensor) -> bool: return x.device.type == "meta" self.assertTrue(tree_all_only(torch.Tensor, is_meta_device, r)) instantiate_device_type_tests(TestDeviceUtils, globals()) class TestCppExtensionUtils(TestCase): def test_cpp_compiler_is_ok(self): self.assertTrue(torch.utils.cpp_extension.check_compiler_ok_for_platform("c++")) def test_cc_compiler_is_ok(self): self.assertTrue(torch.utils.cpp_extension.check_compiler_ok_for_platform("cc")) class TestTraceback(TestCase): def test_basic(self): source = """\ def f(x): def g(x): raise RuntimeError # HEYA x = x * 3 return g(x) + 1 """ out: Dict[str, Any] = {} scope = {"__compile_source__": source} exec(source, scope, out) try: with report_compile_source_on_error(): out["f"](1) except RuntimeError as e: self.assertIn("HEYA", "".join(traceback.format_tb(e.__traceback__))) def test_format_traceback_short(self): try: raise RuntimeError except RuntimeError as e: self.assertRegex( format_traceback_short(e.__traceback__), r".*test_utils.py:\d+ in test_format_traceback_short", ) def test_captured_traceback(self): self.assertIn( "test_captured_traceback", "".join(CapturedTraceback.extract().format()) ) def test_captured_traceback_format_all(self): rs = CapturedTraceback.format_all( [CapturedTraceback.extract(), CapturedTraceback.extract()] ) self.assertEqual(len(rs), 2) self.assertIn("test_captured_traceback_format_all", "".join(rs[0])) def test_captured_traceback_format_all_cached(self): tb = CapturedTraceback.extract() tb.format() # cached rs = CapturedTraceback.format_all([tb, CapturedTraceback.extract()]) self.assertEqual(len(rs), 2) self.assertIn("test_captured_traceback_format_all", "".join(rs[0])) if __name__ == "__main__": run_tests()