diff --git a/aidge_export_arm_cortexm/export_registry.py b/aidge_export_arm_cortexm/export_registry.py index 46c0653ad05895d13ca0d9be7fdb5f4314f2f775..a6e60b3addab7d4c4a49bcd12cd5c42f2aea7a90 100644 --- a/aidge_export_arm_cortexm/export_registry.py +++ b/aidge_export_arm_cortexm/export_registry.py @@ -1,13 +1,31 @@ from aidge_core.export_utils import ExportLib +from aidge_export_arm_cortexm.utils import ROOT class ExportLibAidgeARM(ExportLib): _name="aidge_arm" + def __init__(self, operator): + super(ExportLibAidgeARM, self).__init__(operator) + self.forward_template = str(ROOT / "_Aidge_Arm" / "templates" / "forward_call" / "forward.jinja") + # TODO ugly fix for Tensor registration issue... import aidge_core aidge_core.register_Tensor(["aidge_arm", aidge_core.dtype.float32], aidge_core.get_key_value_Tensor(["cpu", aidge_core.dtype.float32])) +aidge_core.register_Tensor(["aidge_arm", aidge_core.dtype.int32], + aidge_core.get_key_value_Tensor(["cpu", aidge_core.dtype.int32])) +aidge_core.register_Tensor(["aidge_arm", aidge_core.dtype.int8], + aidge_core.get_key_value_Tensor(["cpu", aidge_core.dtype.int8])) +aidge_core.register_Tensor(["aidge_arm", aidge_core.dtype.int4], + aidge_core.get_key_value_Tensor(["cpu", aidge_core.dtype.int4])) +aidge_core.register_Tensor(["aidge_arm", aidge_core.dtype.uint4], + aidge_core.get_key_value_Tensor(["cpu", aidge_core.dtype.uint4])) +aidge_core.register_Tensor(["aidge_arm", aidge_core.dtype.dual_int4], + aidge_core.get_key_value_Tensor(["cpu", aidge_core.dtype.dual_int4])) +aidge_core.register_Tensor(["aidge_arm", aidge_core.dtype.dual_uint4], + aidge_core.get_key_value_Tensor(["cpu", aidge_core.dtype.dual_uint4])) class ExportLibCMSISNN(ExportLib): _name="export_cmsisnn" +