From a2f5fbf09da87783765b5276b04e88749e1980b6 Mon Sep 17 00:00:00 2001 From: peterjc123 Date: Wed, 20 Sep 2023 14:54:10 +0000 Subject: [PATCH] [converter] fix elementwise pass for pack ops --- tinynn/converter/operators/optimize.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tinynn/converter/operators/optimize.py b/tinynn/converter/operators/optimize.py index bdea53cf..63e67086 100644 --- a/tinynn/converter/operators/optimize.py +++ b/tinynn/converter/operators/optimize.py @@ -1811,10 +1811,14 @@ def elementwise_op_transpose_passthrough_pass(self, quantizable_ops_only: bool = tensor_node_dict[op_out.name] = self.graph.graph.vs.find(name=self.graph.tensor_node_map[op_out.name]) # OP specific dim handling logic - if node['node_type'] in (ExtendedOperator.CONCATENATION, ExtendedOperator.GATHER): + if node['node_type'] in (ExtendedOperator.CONCATENATION, ExtendedOperator.GATHER, ExtendedOperator.UNPACK): old_axis = op.axis new_axis = np.where(inv_perm_arr == old_axis)[0][0] op.axis = new_axis + elif node['node_type'] == ExtendedOperator.PACK: + old_axis = op.axis + new_axis = np.where(inv_perm_arr_post == old_axis)[0][0] + op.axis = new_axis elif node['node_type'] == ExtendedOperator.SPLIT_V: old_dim = op.inputs[2].tensor new_dim = np.where(inv_perm_arr == old_dim)[0][0]