From 9c3153acd6d349cde598da49b43df1b9259f5ccf Mon Sep 17 00:00:00 2001 From: Oliver Rausch Date: Fri, 7 Aug 2020 23:05:20 +0200 Subject: [PATCH] Improve shape inference for OneHot (#4452) * Improve shape inference for OneHot Attempt to get the depth parameter before adding a new symbolic dimension. * Update symbolic shape infer * Nit --- .../providers/nuphar/scripts/symbolic_shape_infer.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/nuphar/scripts/symbolic_shape_infer.py b/onnxruntime/core/providers/nuphar/scripts/symbolic_shape_infer.py index cc33355e11..cc1e730fe4 100755 --- a/onnxruntime/core/providers/nuphar/scripts/symbolic_shape_infer.py +++ b/onnxruntime/core/providers/nuphar/scripts/symbolic_shape_infer.py @@ -776,10 +776,13 @@ class SymbolicShapeInference: vi.CopyFrom(helper.make_tensor_value_info(node.output[0], vi.type.tensor_type.elem_type, [input_rank, nz_len])) def _infer_OneHot(self, node): - shape = self._get_shape(node, 0) + sympy_shape = self._get_sympy_shape(node, 0) + depth = self._try_get_value(node, 1) axis = get_attribute(node, 'axis', -1) - axis = handle_negative_axis(axis, len(shape)+1) - new_shape = shape[:axis] + [self._new_symbolic_dim_from_output(node)] + shape[axis:] + axis = handle_negative_axis(axis, len(sympy_shape) + 1) + new_shape = get_shape_from_sympy_shape( + sympy_shape[:axis] + [self._new_symbolic_dim_from_output(node) if not is_literal(depth) else depth] + + sympy_shape[axis:]) vi = self.known_vi_[node.output[0]] vi.CopyFrom(helper.make_tensor_value_info(node.output[0], self.known_vi_[node.input[2]].type.tensor_type.elem_type, new_shape))