Repository navigation
fix(patch): keep *_like factories out of _device_constructors() - #117
Merged
yeahdongcn merged 2 commits intoSep 28, 2026
Merged
yeahdongcn merged 2 commits into
yeahdongcn merged 2 commits into
Conversation
popsiclexu
force-pushed
the
fix/like-factories-device-injection
branch
from
September 24, 2026 11:14
abc9a38 to
a6eb83f
Compare
popsiclexu
marked this pull request as ready for review
September 24, 2026 11:22
Collaborator
|
Confirmed: this behavior differs from CUDA. >>> import torch
>>> src = torch.zeros(4, device="cuda")
>>> torch.set_default_device('cpu')
>>> torch.zeros_like(src).device
device(type='cuda', index=0)
>>> |
yeahdongcn
approved these changes
Sep 28, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Bug
After
import torchada, the*_likefactories stop inheriting the device of their input tensoras soon as a device context is active:
The same applies to
ones_like,empty_like,full_like,rand_like,randn_likeandrandint_like. Note thattorch.set_default_device(d)is not a no-op even whendequals thecurrent device — pushing the context is enough to trigger it.
This is not a wrong argument being passed somewhere, it is a silent change of the result: the
returned tensor owns host memory instead of device memory. Code that then uses it on the device
(copying it into a device buffer, passing it as a kernel argument, ...) dereferences a host
pointer; the outcome is an illegal memory access whose aftermath is usually a stalled context
rather than a clean exception, which makes it hard to attribute to torchada.
Root cause
apply_patches()wraps the torch tensor factories so that an explicitdevice="cuda"istranslated to
"musa", and then re-adds the unwrapped originals to torch's device-injectionregistry, because
__torch_function__dispatch hands out the original C functions and not thePython wrappers we installed:
That registry,
torch.utils._device._device_constructors(), is exactly the set thatDeviceContext.__torch_function__consults to injectdevice=into a call where the caller didnot pass one:
torch deliberately keeps the
*_likefamily out of it, because*_likemust derive its devicefrom the input tensor. Adding them back made torchada override that inheritance, so any active
device context — including a
cpuone — was injected intozeros_like& friends.Wrapping the factories was correct; only their membership in that registry was wrong.
Fix
_discover_factory_functions()now returns two lists: the factories to wrap, and the subset ofthem that torch itself device-injects into (read straight out of torch's own registry rather than
from a hand-maintained list). Only that subset is added back:
The
*_likefactories are still wrapped, so explicittorch.zeros_like(x, device="cuda")->"musa"translation, the JIT builtin aliasing and the AOT-cache handling are unchanged; they aresimply no longer registered for device injection. The fallback path (if the private registry
ever disappears) drops
*_likeby name.The registry ends up at torch's own size again: 45 -> 36 on both torch 2.9.0 and torch 2.11.0.
The 9 entries no longer added are the 7
*_likefactories plusfrom_fileandnormal, none ofwhich torch device-injects.
Tests
Three regression tests in
tests/test_device_strings.py:test_like_factories_are_not_device_injected— the unwrapped*_likefunctions must not be in_device_constructors(), whileempty/zeros/onesmust still be there.test_like_factories_keep_input_device_in_device_context— insidewith torch.device("cpu"),zeros_like/ones_like/empty_like/full_likeof a MUSA tensor must stay on MUSA.test_like_factories_still_translate_an_explicit_cuda_device— the family must stay wrapped, soan explicitly passed
device="cuda"still lands on MUSA:torch.ones_like(x, device="cuda")has
device.type == "musa", also inside a device context. (This is a guard against dropping thefamily from the wrap list; it passes before and after the fix.)
The first two fail before the change and pass after it; the third passes before and after (it is a
guard, not a reproduction).
Full suite,
pytest -q tests/, on MUSA:Nothing regressed on either version -- apart from the new tests' state change, both runs are
stable. The remaining 2.9.0 failures are pre-existing and unrelated to this change
(
tests/test_mm_out_dtype.py::TestMMOutDtypeContract, unchanged by this PR), and on 2.11.0 thesuite is fully green.