saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix shape inference issue in Gather op (#9147)

Browse source 
* add initializer checker for Gather with 1D input

* Check if indices value exists

* Update symbolic_shape_infer.py

* add unit test

* Update symbolic_shape_infer.py

* Update symbolic_shape_infer.py
This commit is contained in:
stevenlix 2021-09-28 22:46:12 -07:00 committed by GitHub
parent b606005858
commit 4f10024868
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 31 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
onnxruntime/python/tools/symbolic_shape_infer.py
View file

@ -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)

21
onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py
View file

@ -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 = [