You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
13 lines
412 B
13 lines
412 B
4 days ago
|
# mypy: allow-untyped-defs
|
||
|
|
import torch
|
||
|
|
from torch._ops import OpOverload, OpOverloadPacket
|
||
|
|
|
||
|
|
|
||
|
|
def _register_decomposition(op: OpOverload, graph: torch._C.Graph) -> None:
|
||
|
|
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)
|