|
import inspect |
|
from typing import Dict, List, Union |
|
|
|
from torch import _C |
|
from torch.onnx import _constants |
|
from torch.onnx._internal import registration |
|
|
|
|
|
class _TorchSchema: |
|
def __init__(self, schema: Union[_C.FunctionSchema, str]) -> None: |
|
if isinstance(schema, _C.FunctionSchema): |
|
self.name: str = schema.name |
|
self.overload_name: str = schema.overload_name |
|
self.arguments: List[str] = [arg.name for arg in schema.arguments] |
|
self.optional_arguments: List[str] = [] |
|
self.returns: List[str] = [ret.name for ret in schema.returns] |
|
self.opsets: List[int] = [] |
|
else: |
|
self.name = schema |
|
self.overload_name = "" |
|
self.arguments = [] |
|
self.optional_arguments = [] |
|
self.returns = [] |
|
self.opsets = [] |
|
|
|
def __str__(self) -> str: |
|
s = ( |
|
f"{self.name}.{self.overload_name}(" |
|
+ ", ".join(self.arguments) |
|
+ ") -> (" |
|
+ ", ".join(self.returns) |
|
+ ")" |
|
+ " in opsets " |
|
+ ", ".join(str(opset) for opset in self.opsets) |
|
) |
|
return s |
|
|
|
def __hash__(self): |
|
|
|
return hash(self.name) |
|
|
|
def __eq__(self, other) -> bool: |
|
if not isinstance(other, _TorchSchema): |
|
return False |
|
|
|
return self.name == other.name |
|
|
|
def is_aten(self) -> bool: |
|
return self.name.startswith("aten::") |
|
|
|
def is_backward(self) -> bool: |
|
return "backward" in self.name |
|
|
|
|
|
def _symbolic_argument_count(func): |
|
params = [] |
|
signature = inspect.signature(func) |
|
optional_params = [] |
|
for name, parameter in signature.parameters.items(): |
|
if name in {"_outputs", "g"}: |
|
continue |
|
if parameter.default is parameter.empty: |
|
optional_params.append(parameter) |
|
else: |
|
params.append(str(parameter)) |
|
return params |
|
|
|
|
|
def all_forward_schemas() -> Dict[str, _TorchSchema]: |
|
"""Returns schemas for all TorchScript forward ops.""" |
|
torch_schemas = [_TorchSchema(s) for s in _C._jit_get_all_schemas()] |
|
return {schema.name: schema for schema in torch_schemas if not schema.is_backward()} |
|
|
|
|
|
def all_symbolics_schemas() -> Dict[str, _TorchSchema]: |
|
"""Returns schemas for all onnx supported ops.""" |
|
symbolics_schemas = {} |
|
|
|
for name in registration.registry.all_functions(): |
|
func_group = registration.registry.get_function_group(name) |
|
assert func_group is not None |
|
symbolics_schema = _TorchSchema(name) |
|
func = func_group.get(_constants.ONNX_MAX_OPSET) |
|
if func is not None: |
|
symbolics_schema.arguments = _symbolic_argument_count(func) |
|
symbolics_schema.opsets = list( |
|
range(func_group.get_min_supported(), _constants.ONNX_MAX_OPSET + 1) |
|
) |
|
else: |
|
|
|
func = func_group.get(7) |
|
symbolics_schema.arguments = _symbolic_argument_count(func) |
|
symbolics_schema.opsets = list(range(7, _constants.ONNX_BASE_OPSET)) |
|
|
|
symbolics_schemas[name] = symbolics_schema |
|
|
|
return symbolics_schemas |
|
|
