Skip to content
Snippets Groups Projects

Feat/operator squeeze unsqueeze

39 files
+ 577
95
Compare changes
  • Side-by-side
  • Inline
Files
39
"""
Copyright (c) 2024 CEA-List
This program and the accompanying materials are made available under the
terms of the Eclipse Public License 2.0 which is available at
http://www.eclipse.org/legal/epl-2.0.
SPDX-License-Identifier: EPL-2.0
"""
import aidge_core
import onnx
from onnx import helper
from aidge_onnx.node_export import auto_register_export
@auto_register_export("Unsqueeze")
def export_unsqueeze(
aidge_node: aidge_core.Node,
node_inputs_name,
node_outputs_name,
verbose: bool = False) -> None:
aidge_operator = aidge_node.get_operator()
onnx_node = helper.make_node(
name=aidge_node.name(),
op_type="Unsqueeze",
inputs=node_inputs_name,
outputs=node_outputs_name,
)
onnx_node.attribute.append(
helper.make_attribute(
"axis",
aidge_operator.get_attr("Axis")
))
return [onnx_node]
Loading