Skip to content
Snippets Groups Projects
Commit 36cb84ee authored by Cyril Moineau's avatar Cyril Moineau
Browse files

[Add] support for registration of export of GenericOperator.

parent 9026834e
No related branches found
No related tags found
3 merge requests!279v0.4.0,!253v0.4.0,!247Doc export
...@@ -84,7 +84,6 @@ class ExportLib(aidge_core.OperatorImpl): ...@@ -84,7 +84,6 @@ class ExportLib(aidge_core.OperatorImpl):
return export_node return export_node
return None return None
# Decorator to register kernels for this export
@classmethod @classmethod
def register(cls, op_type, spec): def register(cls, op_type, spec):
"""Decorator to register an operator implementation for a specified operator type. """Decorator to register an operator implementation for a specified operator type.
...@@ -127,7 +126,6 @@ class ExportLib(aidge_core.OperatorImpl): ...@@ -127,7 +126,6 @@ class ExportLib(aidge_core.OperatorImpl):
return Wrapper return Wrapper
return decorator return decorator
# Decorator to register kernels for this export
@classmethod @classmethod
def register_metaop(cls, op_type, spec): def register_metaop(cls, op_type, spec):
"""Decorator to register a MetaOperator with the export library. """Decorator to register a MetaOperator with the export library.
...@@ -161,3 +159,38 @@ class ExportLib(aidge_core.OperatorImpl): ...@@ -161,3 +159,38 @@ class ExportLib(aidge_core.OperatorImpl):
spec.attrs.add_attr("type", type_name) # MetaOperator specs need to verify the type spec.attrs.add_attr("type", type_name) # MetaOperator specs need to verify the type
return Wrapper return Wrapper
return decorator return decorator
@classmethod
def register_generic(cls, op_type, spec):
"""Decorator to register a GenericOperator with the export library.
Registers a GenericOperator under a given operator type and specification. This decorator
is intended for operator types that are grouped as meta operators.
:param op_type: Operator type(s) to register as a ``GenericOperator``.
:type op_type: Union[str, List[str]]
:param spec: Implementation specification for the GenericOperator.
:type spec: aidge_core.ImplSpec
:return: A wrapper class that initializes the registered GenericOperator.
:rtype: Callable
"""
def decorator(operator):
class Wrapper(operator):
def __init__(self, *args, **kwargs):
return operator(*args, **kwargs)
type_list = []
if isinstance(op_type, list):
type_list = op_type
elif isinstance(op_type, str):
type_list = [op_type]
else:
raise TypeError("Argument 'op_type' of register method should be of type 'List[str]' or 'str', got {type(type)}")
for type_name in type_list:
if (type_name not in cls._export_node_registry):
cls._export_node_registry[type_name] = []
cls._export_node_registry[type_name].append((spec, operator))
aidge_core.register_GenericOperatorOp([cls._name, type_name], cls)
spec.attrs.add_attr("type", type_name) # GenericOperator specs need to verify the type
return Wrapper
return decorator
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment