Keep axis in clone (#13)

* keep axis attribute in inserted onnx dequantize clone * Conditionally add dequant axis
tenstorrent · Aug 6, 2024 · d28d523 · d28d523
1 parent a495355
commit d28d523
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/python/tvm/contrib/pybuda_compile.py b/python/tvm/contrib/pybuda_compile.py
@@ -489,12 +489,13 @@ def duplicate_dequantize_nodes_in_onnx_graph(onnx_module):
             for i, consumer_name in enumerate(consumers):
                 new_node_name = node.name + f"_clone{i}"
                 new_output_name = output_name + f"_clone{i}"
-
+                attrs = {"axis": node.attribute[0].i} if len(node.attribute) > 0 else {}
                 cloned_node = onnx.helper.make_node(
                     node.op_type,
                     node.input,
                     [new_output_name],
-                    name=new_node_name
+                    name=new_node_name,
+                    **attrs
                 )
 
                 # Add the cloned node to the list of nodes to add