pytorch - 💡(How to fix) Fix DISABLED test_vmap_vmap_cpu (__main__.TestComposabilityCPU)

Traceback (most recent call last): File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_utils.py", line 3444, in wrapper method(*args, **kwargs) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_device_type.py", line 459, in instantiated_test raise rte File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_device_type.py", line 439, in instantiated_test result = test(self, **param_kwargs) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 2432, in call result = self._torchdynamo_orig_backend( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 2150, in call result = self._inner_convert( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 735, in call result = _compile( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1919, in _compile guarded_code, tracer_output = compile_inner(code, one_graph, hooks) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_utils_internal.py", line 96, in wrapper_function return function(*args, **kwargs) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1535, in compile_inner result = _compile_inner(code, one_graph, hooks) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1594, in _compile_inner dynamo_output = compile_frame( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1442, in compile_frame bytecode, tracer_output = transform_code_object(code, transform) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/bytecode_transformation.py", line 1626, in transform_code_object tracer_output = transformations(instructions, code_options) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 1414, in transform tracer_output = trace_frame( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 342, in _fn return fn(args, **kwargs) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 875, in trace_frame run_tracer() File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/convert_frame.py", line 856, in run_tracer tracer.run() File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1813, in run while self.step(): File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1480, in step self.dispatch_table[inst.opcode](self, inst) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1017, in wrapper return inner_fn(self, inst) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 2787, in CALL_FUNCTION self.call_function(fn, args, {}) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/symbolic_convert.py", line 1381, in call_function self.push(fn.call_function(self, args, kwargs)) # type: ignore[arg-type] File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/variables/misc.py", line 1341, in call_function return self.obj.call_method(tx, self.name, list(args), kwargs) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/variables/tensor.py", line 935, in call_method result = wrap_fx_proxy(tx, proxy) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 3083, in wrap_fx_proxy return wrap_fx_proxy_cls(target_cls=TensorVariable, **kwargs) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 3158, in wrap_fx_proxy_cls out: VTTypeAlias = _wrap_fx_proxy( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/variables/builder.py", line 3282, in _wrap_fx_proxy example_value = get_fake_value(proxy.node, tx, allow_non_graph_fake=True) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 3736, in get_fake_value return _get_fake_value_impl(node, tx, allow_non_graph_fake) File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 3927, in _get_fake_value_impl _wrap_graph_break_with_torch_runtime_err( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 3725, in _wrap_graph_break_with_torch_runtime_err raise exc.with_traceback(e.traceback) from None File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 3722, in _wrap_graph_break_with_torch_runtime_err gb_fn() File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 3928, in <lambda> lambda: unimplemented( File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/_dynamo/exc.py", line 653, in unimplemented raise Unsupported( torch._dynamo.exc.TorchRuntimeError: RuntimeError when making fake tensor call Explanation: Dynamo failed to run FX node with fake tensors: call_method sin((FakeTensor(..., size=(2, 3)),), **{}): got RuntimeError('tls_on_entry.has_value() INTERNAL ASSERT FAILED at "/__w/pytorch/pytorch/aten/src/ATen/core/PythonFallbackKernel.cpp":53, please report a bug to PyTorch. ') Hint: Your code may result in an error when running in eager. Please double check that your code doesn't contain a similar error when actually running eager/uncompiled. You can do this by removing the torch.compile call, or by using torch.compiler.set_stance("force_eager").

Developer debug context:

For more details about this graph break, please visit: https://meta-pytorch.github.io/compile-graph-break-site/gb/gb4315.html

from user code: File "/__w/pytorch/pytorch/test/functorch/test_eager_transforms.py", line 3590, in test_vmap_vmap self.assertEqual(y, x.sin())

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

To execute this test, run the following from the base repo dir: PYTORCH_TEST_WITH_DYNAMO=1 python test/functorch/test_eager_transforms.py TestComposabilityCPU.test_vmap_vmap_cpu

This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0

TL;DR

The most likely fix for the failing test is to modify the test code to handle the error when running in eager mode or to remove the torch.compile call.

Guidance

• Review the test code in functorch/test_eager_transforms.py to ensure it doesn't contain errors when running in eager mode.
• Try removing the torch.compile call or using torch.compiler.set_stance("force_eager") to see if the test passes.
• Set TORCHDYNAMO_VERBOSE=1 and TORCH_LOGS="+dynamo" to get more detailed logs and internal stack traces.
• Check the graph break site at https://meta-pytorch.github.io/compile-graph-break-site/gb/gb4315.html for more information on the error.

Example

No specific code example is provided, but the test code in functorch/test_eager_transforms.py should be reviewed and modified as needed.

Notes

The error message suggests that the test code may be causing an error when running in eager mode, and the torch.compile call may be contributing to the issue. By removing the torch.compile call or using torch.compiler.set_stance("force_eager"), the test may pass.

Recommendation

Apply a workaround by modifying the test code to handle the error when running in eager mode, or remove the torch.compile call to see if the test passes. This is because the error message suggests that the test code may be causing an error when running in eager mode, and modifying the test code or removing the torch.compile call may resolve the issue.