Skip to content

fix(patch): keep *_like factories out of _device_constructors() - #117

Merged
yeahdongcn merged 2 commits into
MooreThreads:mainfrom
popsiclexu:fix/like-factories-device-injection
Sep 28, 2026
Merged

yeahdongcn merged 2 commits into
MooreThreads:mainfrom
popsiclexu:fix/like-factories-device-injection

Conversation

@popsiclexu

@popsiclexu popsiclexu commented Sep 24, 2026 •

Copy link
Copy Markdown
Contributor

Bug

After import torchada, the *_like factories stop inheriting the device of their input tensor
as soon as a device context is active:

import torch, torchada

src = torch.zeros(4, device="musa")

with torch.device("cpu"):                 # torch.set_default_device("cpu") does the same
    print(torch.zeros_like(src).device)
    # expected: musa:0
    # actual:   cpu          <- before this PR

The same applies to ones_like, empty_like, full_like, rand_like, randn_like and
randint_like. Note that torch.set_default_device(d) is not a no-op even when d equals the
current 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 explicit device="cuda" is
translated to "musa", and then re-adds the unwrapped originals to torch's device-injection
registry, because __torch_function__ dispatch hands out the original C functions and not the
Python wrappers we installed:

for orig_fn in original_fns:       # every wrapped factory, including the *_like family
    constructors.add(orig_fn)

That registry, torch.utils._device._device_constructors(), is exactly the set that
DeviceContext.__torch_function__ consults to inject device= into a call where the caller did
not pass one:

if func in _device_constructors() and kwargs.get("device") is None:
    kwargs["device"] = self.device

torch deliberately keeps the *_like family out of it, because *_like must derive its device
from the input tensor. Adding them back made torchada override that inheritance, so any active
device context — including a cpu one — was injected into zeros_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 of
them 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:

for orig_fn in device_injectable_fns:   # no *_like
    constructors.add(orig_fn)

The *_like factories are still wrapped, so explicit torch.zeros_like(x, device="cuda") ->
"musa" translation, the JIT builtin aliasing and the AOT-cache handling are unchanged; they are
simply no longer registered for device injection. The fallback path (if the private registry
ever disappears) drops *_like by 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 *_like factories plus from_file and normal, none of
which torch device-injects.

Tests

Three regression tests in tests/test_device_strings.py:

  • test_like_factories_are_not_device_injected — the unwrapped *_like functions must not be in
    _device_constructors(), while empty/zeros/ones must still be there.
  • test_like_factories_keep_input_device_in_device_context — inside with torch.device("cpu"),
    zeros_like/ones_like/empty_like/full_like of a MUSA tensor must stay on MUSA.
  • test_like_factories_still_translate_an_explicit_cuda_device — the family must stay wrapped, so
    an 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 the
    family 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:

torch failed passed skipped
2.9.0 before 8 589 18
2.9.0 after 6 591 18
2.11.0 before 2 594 19
2.11.0 after 0 596 19

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 the
suite is fully green.

@popsiclexu
popsiclexu force-pushed the fix/like-factories-device-injection branch from abc9a38 to a6eb83f Compare September 24, 2026 11:14
@popsiclexu
popsiclexu marked this pull request as ready for review September 24, 2026 11:22
@yeahdongcn

Copy link
Copy Markdown
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
yeahdongcn merged commit b3c13ca into MooreThreads:main Sep 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants