This PR refactors `test/test_python_dispatch.py` test file to separate device-related tests from device-unrelated ones.
### Review notes
I split changes into two main commits:
| commit | description |
| - | - |
| **Only change functions order.** | This commit only changes order of test functions. **NO FUNCTIONAL CHANGES AT ALL** |
| **Refactor.** | This commit applies the real refactor. |
Any further commits are addressing review comments.
### Changes
- Extract cuda-specific `test_override_cuda_with_jiterator` from `TestPythonRegistration` into newly created cuda-specific `TestPythonRegistrationCUDA`.
- Remove `test_override_cuda_with_jiterator` function-local `torch.cuda.is_available()` check and apply it as whole `TestPythonRegistrationCUDA` class skip using `@unittest.skipUnless(TEST_CUDA and not TEST_WITH_ROCM, "CUDA-only tests")`.
- Add `HardwareClassification` to all test classes in the file.
- Add `Device` suffix to device-generic `TestWrapperSubclassAliasing` class making it `TestWrapperSubclassAliasingDevice`.
- Propagate rename changes to skips from `test/dynamo_skips` as suggested in review.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/194892
Approved by: https://github.com/pbielak, https://github.com/etaf, https://github.com/jansel
Library.define() builds the op qualname as `self.ns + "::" + name`, where name
is the part of the schema before "(". The C++ schema parser accepts a name that
is already prefixed with the library's own namespace (e.g.
Library("mylib").define("mylib::foo(...) -> ...")) and registers the op
correctly, but the Python side then double-prepends, recording
"mylib::mylib::foo" in _op_defs and the module-global _defs.
That malformed qualname later breaks teardown: _clear_torch_ops_cache does
`ns, name = qualname.split("::")`, which raises "ValueError: too many values to
unpack (expected 2)". Because the cleanup runs from a weakref.finalize
finalizer, in practice this surfaces as an "Exception ignored in" traceback at
interpreter shutdown -- benign (the process still exits 0) but noisy, and it
pollutes CI failure classification (observed inflating "flaky" signal on
otherwise-green jobs).
Root-cause fix: normalize in define -- only prepend the namespace when name does
not already contain "::", mirroring the existing pattern in
register_symm_mem_args. This keeps _op_defs/_defs entries well-formed.
An assert/raise in define was considered and rejected: the genuinely-wrong case
(a namespace that does not match the library) is already caught by C++ with a
clear, actionable error, and the matching-namespace case is one C++ deliberately
accepts -- so rejecting it in Python would be an inconsistent behavior change.
Defense-in-depth: also make _clear_torch_ops_cache skip any qualname that does
not split into exactly "namespace::name" rather than unpacking blindly, since it
runs in a shutdown-time finalizer that must never raise.
Test Plan:
```
# The normalize test fails before the define fix (records "ns::ns::my_op"):
python test/test_python_dispatch.py TestPythonRegistration.test_define_normalizes_namespaced_schema_name
# Full class after the fix: 30 passed
python test/test_python_dispatch.py TestPythonRegistration
lintrunner -a torch/library.py test/test_python_dispatch.py
```
Authored with Claude.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/187970
Approved by: https://github.com/frgossen
Found the following error when vLLM engine exist when upgrade torch after 20260501 nightly version. Introduced via #181785
```
Traceback (most recent call last):
File "/usr/lib/python3.12/weakref.py", line 666, in _exitfunc
f()
File "/usr/lib/python3.12/weakref.py", line 590, in __call__
return info.func(*info.args, **(info.kwargs or {}))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/venv/lib/python3.12/site-packages/torch/library.py", line 628, in _del_library
_clear_torch_ops_cache(op_defs)
File "/opt/venv/lib/python3.12/site-packages/torch/library.py", line 593, in _clear_torch_ops_cache
if not hasattr(namespace, name):
^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/venv/lib/python3.12/site-packages/torch/_ops.py", line 1383, in __getattr__
op, overload_names = _get_packet(qualified_op_name, module_name)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/venv/lib/python3.12/site-packages/torch/_ops.py", line 1408, in _get_packet
op, overload_names = torch._C._jit_get_operation(qualname)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
UnicodeDecodeError: 'utf-8' codec can't decode byte 0xe2 in position 275: invalid continuation byte
```
The issue is that `hasattr(namespace, name)` is not safe to call inside a weakref finalizer during interpreter shutdown when name is not already cached on the namespace object. If the op was already cached (via a prior getattr that set it as an attribute), hasattr would find it directly without calling into C++. But if it wasn't cached, it falls through to `__getattr__` which calls into the C++ layer that may already be torn down.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/186611
Approved by: https://github.com/zou3519
Make 3 tests in TestWrapperSubclassAliasing properly use the device parameter from instantiate_device_type_tests, enabling them to run on all accelerators (CPU, CUDA, XPU, MPS).
Modified tests:
- test_wrapper_subclass_aliasing_conv2d
- test_wrapper_subclass_aliasing_out_op
- test_wrapper_subclass_aliasing_fft_fft2
Test plan:
TEST_CONFIG=cuda python test/test_python_dispatch.py TestWrapperSubclassAliasingCUDA -v
TEST_CONFIG=cpu python test/test_python_dispatch.py TestWrapperSubclassAliasingCPU -v
All 20 tests pass on both CPU and CUDA.
CC: @fffrog @albanD @jbschlosser
Pull Request resolved: https://github.com/pytorch/pytorch/pull/185068
Approved by: https://github.com/Skylion007
Summary:
Change the default behavior of `torch.library.Library` routines to allow
multiple overrides of the same symbol without error by default, instead
of requiring explicit opt-in. This means the python implementation of
`Library` now matches the underlying C++ implementation.
Done via. changing the default for the `allow_override` argument from
`False` -> `True` on all relevant methods.
With the addition of op overrides via. `torch._native.ops`, we now have
the case where `import torch` defines overrides on `aten::` symbols,
which downstream libraries may also attempt to override. With the
previous default this will cause an exception, a BC-breaking change.
Notably seen in https://github.com/pytorch/pytorch/issues/181122
Discussed with albanD, malfet, drisspg, ngimel, zou3519 and others.
All relevant tests updated to reflect the new behavior
Test Plan:
```
pytest -sv test/test_custom_ops.py
pytest -sv test/test_python_dispatch.py
pytest -sv test/test_opaque_obj_v2.py
pytest -sv test/test_fake_tensor.py
```
Signed-off-by: Simon Layton <simonlayton@meta.com>
Signed-off-by: Simon Layton <simonlayton@meta.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/182293
Approved by: https://github.com/malfet
Summary:
Change the default behavior of `torch.library.Library` routines to allow
multiple overrides of the same symbol without error by default, instead
of requiring explicit opt-in. This means the python implementation of
`Library` now matches the underlying C++ implementation.
Done via. changing the default for the `allow_override` argument from
`False` -> `True` on all relevant methods.
With the addition of op overrides via. `torch._native.ops`, we now have
the case where `import torch` defines overrides on `aten::` symbols,
which downstream libraries may also attempt to override. With the
previous default this will cause an exception, a BC-breaking change.
Notably seen in https://github.com/pytorch/pytorch/issues/181122
Discussed with albanD, malfet, drisspg, ngimel, zou3519 and others.
All relevant tests updated to reflect the new behavior
Test Plan:
```
pytest -sv test/test_custom_ops.py
pytest -sv test/test_python_dispatch.py
pytest -sv test/test_opaque_obj_v2.py
pytest -sv test/test_fake_tensor.py
```
Signed-off-by: Simon Layton <simonlayton@meta.com>
Signed-off-by: Simon Layton <simonlayton@meta.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/182293
Approved by: https://github.com/malfet
ghstack dependencies: #182428
`my_lib` in `test_storage_preserve_nonhermetic_in_hermetic_context` leaks into global op space after the test ends and affect subsequent tests in the same process using dynamo.
Without the fix, running any tests requiring checkpoint/compile or dynamo-related after `test_storage_preserve_nonhermetic_in_hermetic_context` fails with
```
torch._dynamo.exc.BackendCompilerFailed: backend='aot_eager' raised:
TypeError: 'CustomDecompTable' object is not a mapping
```
e.g. `python -m pytest -v pytorch/test/test_torch.py::TestTorch::test_storage_preserve_nonhermetic_in_hermetic_context pytorch/test/test_autograd.py::TestAutograd::test_checkpoint_compile_no_recompile`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/180998
Approved by: https://github.com/albanD, https://github.com/ezyang
Calling `del lib` on a `torch.library.Library` ran the weakref finalizer but
left the cached `OpOverloadPacket` on `torch.ops.<ns>.<name>` in place. The
C++ side had been reset, but `OpOverloadPacket._overload_names` still listed
'default' while `getattr(packet, 'default')` now raised `AttributeError`. Any
later code that walked all CompositeImplicitAutograd ops would hit that --
most visibly, `{**core_aten_decompositions(), **inductor_decompositions}` at
`torch/_inductor/decomposition.py:108`, where CPython's `DICT_UPDATE` opcode
rewrote the `AttributeError` as the misleading
`TypeError: 'CustomDecompTable' object is not a mapping`.
`Library._destroy()` already did this cleanup, so the explicit-teardown path
(e.g. `_scoped_library`) was fine; only the gc-driven finalizer path was
broken. Extracted the cleanup into `_clear_torch_ops_cache()` and called it
from `_del_library()` as well.
This was hitting slow-gradcheck periodic in test_torch.py:
`test_storage_preserve_nonhermetic_in_hermetic_context` creates a DEF-mode
Library, calls its op, then lets it go out of scope; subsequent CUDA tests
then crashed in `before_cuda_memory_leak_check` -> `torch._dynamo.reset()` ->
the lazy import chain above.
Fixes#181765. The flaky-test disabler had been auto-creating per-test
issues (#181343, #181537, #181684, #119515) papering over this -- those
tests should be re-enabled once this lands.
Authored with Claude.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/181785
Approved by: https://github.com/Lucaskabela
Calling `del lib` on a `torch.library.Library` ran the weakref finalizer but
left the cached `OpOverloadPacket` on `torch.ops.<ns>.<name>` in place. The
C++ side had been reset, but `OpOverloadPacket._overload_names` still listed
'default' while `getattr(packet, 'default')` now raised `AttributeError`. Any
later code that walked all CompositeImplicitAutograd ops would hit that --
most visibly, `{**core_aten_decompositions(), **inductor_decompositions}` at
`torch/_inductor/decomposition.py:108`, where CPython's `DICT_UPDATE` opcode
rewrote the `AttributeError` as the misleading
`TypeError: 'CustomDecompTable' object is not a mapping`.
`Library._destroy()` already did this cleanup, so the explicit-teardown path
(e.g. `_scoped_library`) was fine; only the gc-driven finalizer path was
broken. Extracted the cleanup into `_clear_torch_ops_cache()` and called it
from `_del_library()` as well.
This was hitting slow-gradcheck periodic in test_torch.py:
`test_storage_preserve_nonhermetic_in_hermetic_context` creates a DEF-mode
Library, calls its op, then lets it go out of scope; subsequent CUDA tests
then crashed in `before_cuda_memory_leak_check` -> `torch._dynamo.reset()` ->
the lazy import chain above.
Fixes#181765. The flaky-test disabler had been auto-creating per-test
issues (#181343, #181537, #181684, #119515) papering over this -- those
tests should be re-enabled once this lands.
Authored with Claude.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/181785
Approved by: https://github.com/Lucaskabela
## Summary
- define `TraceableWrapperSubclass` in `torch/utils/_python_dispatch.py` as the canonical traceable wrapper subclass protocol, including the flatten/unflatten invariants and supported `__tensor_unflatten__` forms
- thread that protocol through the fake, AOT, and proxy helper boundaries so the contract is referenced directly instead of living in casts and comments
- add a focused runtime test that covers both `@staticmethod` and `@classmethod` `__tensor_unflatten__` implementations
## Root cause
The traceable wrapper subclass contract was only implied by repeated `hasattr(..., "__tensor_flatten__")` checks and an internal structural type with inaccurate method signatures. That left the canonical meaning of the protocol split across helpers, comments, and downstream assumptions.
## Proposed fix
Introduce an explicit `TraceableWrapperSubclass` protocol in `torch/utils/_python_dispatch.py`, move the contract documentation there, make `is_traceable_wrapper_subclass()` reference that canonical protocol, and update the affected fake/proxy/AOT helpers to use the protocol type directly.
## Why this is the right long term fix
This keeps the runtime behavior unchanged while giving PT2 wrapper-subclass support one documented source of truth. Future call sites can reference the same protocol instead of re-expressing the contract, which reduces drift in both signatures and invariants.
## Testing
- `python3 -m compileall torch/utils/_python_dispatch.py torch/_subclasses/fake_tensor.py torch/_functorch/_aot_autograd/subclass_utils.py torch/fx/experimental/proxy_tensor.py test/test_python_dispatch.py`
- `PYTHONPATH=/tmp/repos/pytorch/pytorch python3 test/test_python_dispatch.py TestPythonDispatch.test_traceable_wrapper_subclass_protocol_runtime_check TestPythonDispatch.test_make_fx_with_subclass TestPythonDispatch.test_make_wrapper_subclass_propagates_metadata`
- `PYTHONPATH=/tmp/repos/pytorch/pytorch python3 test/dynamo/test_subclasses.py SubclassTests.test_deferred_init_subclass_init_not_traced`
Drafted via Codex, published after manual review by @Lucaskabela
Pull Request resolved: https://github.com/pytorch/pytorch/pull/180359
Approved by: https://github.com/Skylion007
Co-authored-by: PyTorch MergeBot <pytorchmergebot@users.noreply.github.com>
fixes https://github.com/pytorch/pytorch/issues/145108
previously we only remove the functional bit in the destructor and we left the backend bit, this PR uses save/restore to save a snapshot of the TLS upon construction and restores that exactly upon destruction. RAII guards don't respect non lexical TLS modification so if someone adds something after construction it will be overwritten when the state is restored
uncommented flakey test and added more unit tests to confirm TLS state is properly restored.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/178536
Approved by: https://github.com/zou3519, https://github.com/ezyang
[relanding again after fixing internal build]
Summary:
This might cause some new DDEs on call sites that do not use is_contiguous_or_false() or sym_is_contiguous()
but want to find those call sites to handle this properly by calling is_contiguous_or_false() and not is_contiguous() explitly when appropriate.
I had to fix one issue after removing the implicit size oblivious reasoning. here is context
we defined in this https://github.com/pytorch/pytorch/pull/157472 sym_is_contiguous to be the function computing contiguity for dynamic shapes in c++. It returns a symbolic expression that represents contiguity and guaranteed not to throw a DDE.
when people call is_contiguous we do sym_is_contiguous().guard_bool()
when people call is_contiguous_or_false we do sym_is_contiguous().guard_or_false()
one issue not handled well was this path
```
c10::SymBool TensorImpl::sym_is_contiguous_custom(
at::MemoryFormat memory_format) const {
if (C10_UNLIKELY(matches_python_custom(SizesStridesPolicy::CustomStrides))) {
return pyobj_slot_.load_pyobj_interpreter()->is_contiguous(
this, memory_format);
}
return sym_is_contiguous_default(memory_format);
}
```
namely if we call sym_is_contiguous_custom but we have matches_python_custom(SizesStridesPolicy::CustomStrides) return true , then we used to call is_contiguous(this, memory_format);
This used to go through the load_pyobj_interpreter and end up calling the python is_contiguous call which used implicit size oblivious reasoning.
once we removed that implicit size oblivious reasoning, the right thing we want is to call
return pyobj_slot_.load_pyobj_interpreter()->sym_is_contiguous(this, memory_format);
otherwise we would get DDE even if the caller is doing sym_is_contiguous.
so I had to define it for pyinterpreter, and then I had to override it for nested tensors.
Approved by: https://github.com/ezyang
Test Plan:
contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/e444cd24d48b3a46f067974f2cc157f5ed27709f
Rollback Plan:
Differential Revision: D80435179
Pull Request resolved: https://github.com/pytorch/pytorch/pull/160869
Approved by: https://github.com/ezyang
Per title, previously we started throwing noisy warnings, but given how popular this pattern was in our test suite decided to leave it as warning, not as silent behavior change for one release.
Now `treatSequenceAsTuple` would return `true` in the only case where the sequence was indeed a tuple, so no need for a special function anymore.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/160794
Approved by: https://github.com/albanD
This might cause some new DDEs on call sites that do not use is_contiguous_or_false() or sym_is_contiguous()
but want to find those call sites to handle this properly by calling is_contiguous_or_false() and not is_contiguous() explitly when appropriate.
I had to fix one issue after removing the implicit size oblivious reasoning. here is context
we defined in this https://github.com/pytorch/pytorch/pull/157472 sym_is_contiguous to be the function computing contiguity for dynamic shapes in c++. It returns a symbolic expression that represents contiguity and guaranteed not to throw a DDE.
when people call is_contiguous we do sym_is_contiguous().guard_bool()
when people call is_contiguous_or_false we do sym_is_contiguous().guard_or_false()
one issue not handled well was this path
```
c10::SymBool TensorImpl::sym_is_contiguous_custom(
at::MemoryFormat memory_format) const {
if (C10_UNLIKELY(matches_python_custom(SizesStridesPolicy::CustomStrides))) {
return pyobj_slot_.load_pyobj_interpreter()->is_contiguous(
this, memory_format);
}
return sym_is_contiguous_default(memory_format);
}
```
namely if we call sym_is_contiguous_custom but we have matches_python_custom(SizesStridesPolicy::CustomStrides) return true , then we used to call is_contiguous(this, memory_format);
This used to go through the load_pyobj_interpreter and end up calling the python is_contiguous call which used implicit size oblivious reasoning.
once we removed that implicit size oblivious reasoning, the right thing we want is to call
return pyobj_slot_.load_pyobj_interpreter()->sym_is_contiguous(this, memory_format);
otherwise we would get DDE even if the caller is doing sym_is_contiguous.
so I had to define it for pyinterpreter, and then I had to override it for nested tensors.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/159197
Approved by: https://github.com/ezyang
Investigation is a separate issue. For now I want to get the CI back up
and running on the other tests. The problem seems to be that
IncludeDispatchKeyGuard doesn't actually reset the state, which seems
very, very wrong.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/145109
Approved by: https://github.com/williamwen42
Fixes https://github.com/pytorch/pytorch/issues/141120
Registering a fallthrough for a backend correctly alters nonFallthroughKeysPerBackend_[backend_idx]. However, the backend_idx calculation does not take into account the local dispatch key set, which is used to temporarily turn on Meta as a backend. This means that makeFallthrough does not behave exactly as if it was a normal function which redispatched rather than a "fake function" implemented with a key mask.
So e.g. impl::computeDispatchKeySet(ks, nonFallthroughKeysPerBackend_[backend_idx]); will exclude keys like Meta which may be in the TLS include set.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/141581
Approved by: https://github.com/bdhirsh
Fixes#136565
This PR makes the python fallback robust to the case where there are no active modes & no tensors with the Python key. In this case, simply redispatch with the Python key disabled.
This was found when trying to use reentrant dispatch for NJT to get decompositions under `inference_mode()` when the autograd key is disabled.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/136566
Approved by: https://github.com/bdhirsh
Fixes https://github.com/pytorch/pytorch/issues/136064
In the linked repro, this issue was that there was some code like this:
```
# x has dtype torch.float32
def f(x):
y = x.view(torch.float32)
y.copy_(...)
```
Where because `view.dtype` is implemented today to potentially directly return its input, we would end up directly clobbering the proxy for our graph input (replacing its FX proxy value from `arg0_1` to `view_1`). This is not desirable, because we have careful assertions in AOTDispatcher that mutations only ever happen on graph inputs - but this clobbering caused the mutation to appear, from the perspective of the FX graph, like it was happening on a view of the input.
Why is this normally not a problem? Ordinarily, the `ADInplaceOrView` kernel for `view.dtype` will take the output of the view kernel, [and detach() it](https://github.com/pytorch/pytorch/blob/main/tools/autograd/gen_inplace_or_view_type.py#L466) (properly creating a fresh `TensorImpl`).
This does **not** happen, though, if you are executing the kernel from with a `__torch_dispatch__` region: the `ADInplaceOrView` logic has already run above you, so that key will be in the TLS exclude set.
This PR changes eager behavior - at first I considered trying to only change behavior under compile. But this problem isn't technically specific to PT2: if you ever rely on tensor identity from inside of a __torch_dispatch__ call, then we need to make sure the raw `view.dtype` kernel doesn't directly return the input.
I am also making the assumption that "`view.dtype` no-op'ing when the dtype is the same" is not a case worth optimizing in eager mode, and that the overhead of the `TensorImpl` creation is relatively negligible.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/136074
Approved by: https://github.com/Skylion007, https://github.com/ezyang, https://github.com/albanD
ghstack dependencies: #136041
Benedykt Bela