mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
d9f0774f98
- Allow registering custom decompositions - Add easier API for invoking decompositions - Shorten API names (no users yet) I am doing these as one pr because they are fairly short/simple and because github first does not support ghstack yet. cc @Chillee @zou3519 Pull Request resolved: https://github.com/pytorch/pytorch/pull/76252 Approved by: https://github.com/davidberard98
9 lines
360 B
Python
9 lines
360 B
Python
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)
|