pytorch - 💡(How to fix) Fix register_fake to throw error if signature doesn't match op's schema [1 participants]

Official PRs (…)
ON THIS PAGE

Recommended Tools

×6

Utilities matched from this issue’s tags and category — try them while you read without losing context.

GitHub issue graph ai analysis

Paste a GitHub issue URL. We fetch that issue, discover linked issues from bodies/comments/timeline, collect linked pull requests, and produce a structured English report.

The report is written in English Markdown for sharing and archival.

Helpful · Quick feedback

Loading…
GitHub stats
pytorch/pytorch#181851Fetched 2026-04-30 06:18:14
View on GitHub
Comments
0
Participants
1
Timeline
31
Reactions
0
Participants
Timeline (top)
mentioned ×12subscribed ×12labeled ×7

Error Message

E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] fake tensor raised TypeError E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] Traceback (most recent call last): E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1487, in torch_dispatch E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] return self.dispatch(func, types, args, kwargs) E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 2264, in dispatch E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] return self._cached_dispatch_impl(func, types, args, kwargs) E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1635, in _cached_dispatch_impl E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] output = self._dispatch_impl(func, types, args, kwargs) E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 2900, in _dispatch_impl E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] result = maybe_fake_impl(*args, **kwargs) E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] File "/opt/pytorch/pytorch/torch/_library/utils.py", line 23, in call E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] return self.func(*args, **kwargs) E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] ^^^^^^^^^^^^^^^^^^^^^^^^^^ E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] File "/opt/pytorch/pytorch/torch/library.py", line 1499, in inner E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] return func(*args, **kwargs) E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] ^^^^^^^^^^^^^^^^^^^^^ E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] File "/opt/pytorch/pytorch/torch/_library/custom_ops.py", line 664, in fake_impl E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] return self._abstract_fn(*args, **kwargs) E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] TypeError: foo_fake() takes 1 positional argument but 2 were given

Fix Action

Fix / Workaround

Error

E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] fake tensor raised TypeError
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] Traceback (most recent call last):
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1487, in __torch_dispatch__
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return self.dispatch(func, types, args, kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 2264, in dispatch
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return self._cached_dispatch_impl(func, types, args, kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1635, in _cached_dispatch_impl
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     output = self._dispatch_impl(func, types, args, kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 2900, in _dispatch_impl
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     result = maybe_fake_impl(*args, **kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_library/utils.py", line 23, in __call__
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return self.func(*args, **kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/library.py", line 1499, in inner
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return func(*args, **kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_library/custom_ops.py", line 664, in fake_impl
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return self._abstract_fn(*args, **kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] TypeError: foo_fake() takes 1 positional argument but 2 were given

Code Example

import torch
from torch.library import custom_op

# Define a custom op with a real kernel that takes (Tensor, int) -> Tensor
@custom_op("mylib::foo", mutates_args=())
def foo(x: torch.Tensor, n: int) -> torch.Tensor:
    return x.clone().repeat(n)

# Register a fake kernel with a DIFFERENT signature (missing the int arg)
@foo.register_fake
def foo_fake(x: torch.Tensor) -> torch.Tensor:
    return x.clone()

result = foo(torch.randn(3), 2)

@torch.compile(fullgraph=True)
def compiled_fn(x):
    return foo(x, 2)

result2 = compiled_fn(torch.randn(3))

---

E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] fake tensor raised TypeError
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] Traceback (most recent call last):
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1487, in __torch_dispatch__
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return self.dispatch(func, types, args, kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 2264, in dispatch
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return self._cached_dispatch_impl(func, types, args, kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1635, in _cached_dispatch_impl
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     output = self._dispatch_impl(func, types, args, kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 2900, in _dispatch_impl
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     result = maybe_fake_impl(*args, **kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_library/utils.py", line 23, in __call__
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return self.func(*args, **kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/library.py", line 1499, in inner
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return func(*args, **kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_library/custom_ops.py", line 664, in fake_impl
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return self._abstract_fn(*args, **kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] TypeError: foo_fake() takes 1 positional argument but 2 were given
RAW_BUFFERClick to expand / collapse

🚀 The feature, motivation and pitch

register_fake should validate that the fake kernel's signature matches the op's schema at registration time. Currently, a mismatched signature (e.g. missing args) is silently accepted and only fails later at torch.compile

import torch
from torch.library import custom_op

# Define a custom op with a real kernel that takes (Tensor, int) -> Tensor
@custom_op("mylib::foo", mutates_args=())
def foo(x: torch.Tensor, n: int) -> torch.Tensor:
    return x.clone().repeat(n)

# Register a fake kernel with a DIFFERENT signature (missing the int arg)
@foo.register_fake
def foo_fake(x: torch.Tensor) -> torch.Tensor:
    return x.clone()

result = foo(torch.randn(3), 2)

@torch.compile(fullgraph=True)
def compiled_fn(x):
    return foo(x, 2)

result2 = compiled_fn(torch.randn(3))

Error

E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] fake tensor raised TypeError
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] Traceback (most recent call last):
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1487, in __torch_dispatch__
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return self.dispatch(func, types, args, kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 2264, in dispatch
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return self._cached_dispatch_impl(func, types, args, kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 1635, in _cached_dispatch_impl
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     output = self._dispatch_impl(func, types, args, kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py", line 2900, in _dispatch_impl
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     result = maybe_fake_impl(*args, **kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_library/utils.py", line 23, in __call__
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return self.func(*args, **kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/library.py", line 1499, in inner
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return func(*args, **kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]   File "/opt/pytorch/pytorch/torch/_library/custom_ops.py", line 664, in fake_impl
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]     return self._abstract_fn(*args, **kwargs)
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E0429 01:11:38.212000 2100392 /opt/pytorch/pytorch/torch/_subclasses/fake_tensor.py:1489] [0/0] TypeError: foo_fake() takes 1 positional argument but 2 were given

Alternatives

Keep the existing behaviour

Additional context

No response

cc @chauhang @penguinwu @bdhirsh @bobrenjc93 @aorenste

extent analysis

TL;DR

The issue can be resolved by ensuring the fake kernel's signature matches the op's schema at registration time.

Guidance

  • Verify that the foo_fake function has the same signature as the foo function, including the same number and types of arguments.
  • Update the foo_fake function to accept an additional int argument, matching the signature of the foo function.
  • Check the documentation for @custom_op and @foo.register_fake to ensure that the correct syntax and semantics are being used.
  • Test the updated code to ensure that the TypeError is resolved and the foo_fake function is called correctly.

Example

@foo.register_fake
def foo_fake(x: torch.Tensor, n: int) -> torch.Tensor:
    return x.clone().repeat(n)

Notes

The provided code snippet and error message suggest that the issue is due to a mismatch between the signature of the foo function and the foo_fake function. Ensuring that these signatures match should resolve the TypeError.

Recommendation

Apply the workaround by updating the foo_fake function to match the signature of the foo function, as shown in the example above. This should resolve the TypeError and allow the code to run correctly.

Vote matrix · Quick signals

Works
Did the solution work? Tap to confirm.
Easy Fix
Was it a quick fix?
Time Saver
Did it save you time?
Blocking
Was it severely blocking?
Common Issue
Are others likely hitting this too?
Flaky / Intermittent
Is it intermittent?
Verified / Reproducible
Can you reproduce it reliably?
Loading…

Still need to ship something?

×6

Another batch ranked right after the header list — different links, same matching logic.

Back to top recommendations

TRENDING