diff --git a/test/inductor/test_compile_subprocess.py b/test/inductor/test_compile_subprocess.py index 51aa7b70b9c..bf474bfbf17 100644 --- a/test/inductor/test_compile_subprocess.py +++ b/test/inductor/test_compile_subprocess.py @@ -206,7 +206,8 @@ class TestSubprocess(TestCase): start = time.time() last_report = start - while _AsyncFxCompile._stat_compiled_runs < 4: + while True: + start_stat_compiled_runs = _AsyncFxCompile._stat_compiled_runs # Sleep a bit so we don't drive the CPU unnecessarily. time.sleep(0.25) @@ -219,6 +220,9 @@ class TestSubprocess(TestCase): # Backward pass output.sum().backward() + if _AsyncFxCompile._stat_compiled_runs - start_stat_compiled_runs == 2: + break + # DEBUGGING: Print a periodic message so we know we're still # running... now = time.time() @@ -231,12 +235,12 @@ class TestSubprocess(TestCase): "Test timed out before producing a compiled artifact." ) - self.assertEqual(_AsyncFxCompile._stat_compiled_runs, 4) + self.assertGreater(_AsyncFxCompile._stat_compiled_runs, 1) # Make sure we ran eager at least once. Normally this will be # something like 80. self.assertGreater(_AsyncFxCompile._stat_eager_runs, 0) - self.assertEqual(_AsyncFxCompile._stat_bg_started, 1) - self.assertEqual(_AsyncFxCompile._stat_bg_finished, 1) + self.assertEqual(_AsyncFxCompile._stat_bg_started, 2) + self.assertEqual(_AsyncFxCompile._stat_bg_finished, 2) if RUN_CPU: diff --git a/test/test_sympy_utils.py b/test/test_sympy_utils.py index 5343e2e0a9f..3f5fcee2ed5 100644 --- a/test/test_sympy_utils.py +++ b/test/test_sympy_utils.py @@ -24,6 +24,7 @@ from torch.utils._sympy.functions import ( FloorDiv, Identity, OpaqueUnaryFn_cos, + BitwiseFn_bitwise_and, simple_floordiv_gcd, ) from torch.utils._sympy.interp import sympy_interp @@ -873,6 +874,10 @@ class TestSympyFunctions(TestCase): r = pickle.loads(pickle.dumps(x)) self.assertEqual(x, r) + x = BitwiseFn_bitwise_and(sympy.Symbol("a"), sympy.Symbol("b")) + r = pickle.loads(pickle.dumps(x)) + self.assertEqual(x, r) + class TestSingletonInt(TestCase): def test_basic(self): diff --git a/torch/utils/_sympy/functions.py b/torch/utils/_sympy/functions.py index 2b6c159f5c3..7d3e5c1c7b4 100644 --- a/torch/utils/_sympy/functions.py +++ b/torch/utils/_sympy/functions.py @@ -1411,7 +1411,10 @@ def make_opaque_bitwise_fn(name, real_op_name): return sympy.Integer(getattr(operator, real_op_name)(int(a), int(b))) return None - BitwiseFn.__name__ = "BitwiseFn_" + name + nm = "BitwiseFn_" + name + BitwiseFn.__name__ = nm + BitwiseFn.__qualname__ = nm + return BitwiseFn