mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
038b927590
See #127836 for details. Pull Request resolved: https://github.com/pytorch/pytorch/pull/127844 Approved by: https://github.com/oulgen ghstack dependencies: #127842, #127843
13 lines
402 B
Python
13 lines
402 B
Python
# mypy: allow-untyped-defs
|
|
import torch
|
|
from torch._ops import OpOverload, OpOverloadPacket
|
|
|
|
|
|
def _register_decomposition(op: OpOverload, graph: torch._C.Graph):
|
|
assert not isinstance(
|
|
op, OpOverloadPacket
|
|
), f"Must pass specific op overload, not overload packet, found {op}"
|
|
assert isinstance(op, OpOverload)
|
|
|
|
torch._C._jit_register_decomposition_for_schema(op._schema, graph)
|