diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index ebcb0d99e1..e7c7886108 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -893,16 +893,17 @@ class SymbolicShapeInference: data_shape[:axis] + indices_shape + data_shape[axis + 1:])) # for 1D input, do some sympy compute if node.input[0] in self.sympy_data_ and len(data_shape) == 1 and 0 == get_attribute(node, 'axis', 0): - idx = self._get_value(node, 1) - data = self.sympy_data_[node.input[0]] - if type(data) == list: - if type(idx) == np.ndarray and len(idx.shape) == 1: - self.sympy_data_[node.output[0]] = [data[int(i)] for i in idx] + idx = self._try_get_value(node, 1) + if idx is not None: + data = self.sympy_data_[node.input[0]] + if type(data) == list: + if type(idx) == np.ndarray and len(idx.shape) == 1: + self.sympy_data_[node.output[0]] = [data[int(i)] for i in idx] + else: + self.sympy_data_[node.output[0]] = data[int(idx)] else: - self.sympy_data_[node.output[0]] = data[int(idx)] - else: - assert idx == 0 - self.sympy_data_[node.output[0]] = data + assert idx == 0 or idx == -1 + self.sympy_data_[node.output[0]] = data def _infer_GatherElements(self, node): indices_shape = self._get_shape(node, 1) diff --git a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py index 55b50f01fa..e3935c4e40 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py +++ b/onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py @@ -94,6 +94,27 @@ class TestSymbolicShapeInferenceForOperators(unittest.TestCase): ] self._check_shapes(graph, inferred.graph, expected_shapes) + def test_gather_indices(self): + graph = helper.make_graph([ + helper.make_node("Constant", [], ["data"], "constant", + value=helper.make_tensor('input', TensorProto.FLOAT, + [5], [0.0, 1.0, 2.0, 3.0, 4.0])), + helper.make_node("Gather", ["data", "indices"], ["output"], axis=0), + ], "Gather_Test", [ + helper.make_tensor_value_info('indices', TensorProto.INT64, ['b']), + ], [ + helper.make_tensor_value_info('output', TensorProto.FLOAT, ['b']), + ]) + model = helper.make_model(graph, producer_name='Gather_Test_Model') + model.opset_import[0].version = 13 + + inferred = SymbolicShapeInference.infer_shapes(model, auto_merge=True) + expected_shapes = [ + helper.make_tensor_value_info('data', TensorProto.FLOAT, [5]), + helper.make_tensor_value_info('output', TensorProto.FLOAT, ['b']) + ] + self._check_shapes(graph, inferred.graph, expected_shapes) + def test_embed_layer_norm(self): hidden_size = 32 initializers = [