From ddd4586a2fbc047a1404ef78a7bd5063fc87f9ab Mon Sep 17 00:00:00 2001 From: KeDengMS Date: Tue, 24 Aug 2021 16:33:23 -0700 Subject: [PATCH] [Symbolic Shape Infer] add more ops for auto merge (#8824) As Less/Equal/Greater/LessOrEqual/GreaterOrEqual ops can broadcast --- onnxruntime/python/tools/symbolic_shape_infer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py index 0bc5523779..728b0813be 100755 --- a/onnxruntime/python/tools/symbolic_shape_infer.py +++ b/onnxruntime/python/tools/symbolic_shape_infer.py @@ -1797,7 +1797,7 @@ class SymbolicShapeInference: if self.auto_merge_: if node.op_type in [ 'Add', 'Sub', 'Mul', 'Div', 'MatMul', 'MatMulInteger', 'MatMulInteger16', 'Concat', - 'Where', 'Sum' + 'Where', 'Sum', 'Equal', 'Less', 'Greater', 'LessOrEqual', 'GreaterOrEqual' ]: shapes = [self._get_shape(node, i) for i in range(len(node.input))] if node.op_type in ['MatMul', 'MatMulInteger', 'MatMulInteger16']: