Skip to content

Commit 684fc1b

Browse files
shinhlinkerzhang
andcommitted
Keep symbolic dims in Concat with a single input (#2418)
* Keep symbolic dims in Concat with a single input When there is only a single input, we need to discard the symbolic dimension in the specified axis. * Validate axis of Concat with a single input Co-authored-by: Ke Zhang <[email protected]>
1 parent 65020da commit 684fc1b

2 files changed

Lines changed: 12 additions & 0 deletions

File tree

‎onnx/defs/tensor/defs.cc‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -332,6 +332,11 @@ ONNX_OPERATOR_SET_SCHEMA(
332332
axis += rank;
333333
}
334334

335+
if (numInputs == 1) {
336+
propagateShapeFromInputToOutput(ctx, 0, 0);
337+
return;
338+
}
339+
335340
bool all_lengths_known = true;
336341
int total_length = 0;
337342

‎onnx/test/shape_inference_test.py‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,13 @@ def test_concat_param(self): # type: () -> None
199199
[])
200200
self._assert_inferred(graph, [make_tensor_value_info('z', TensorProto.FLOAT, ("a", 5))])
201201

202+
def test_concat_param_single_input(self): # type: () -> None
203+
graph = self._make_graph(
204+
[("x", TensorProto.FLOAT, ("a", 2))],
205+
[make_node("Concat", ['x'], ['z'], axis=0)],
206+
[])
207+
self._assert_inferred(graph, [make_tensor_value_info('z', TensorProto.FLOAT, ("a", 2))])
208+
202209
def test_reshape_dynamic_shape(self): # type: () -> None
203210
graph = self._make_graph(
204211
[('x', TensorProto.UINT8, (2, 4, 3)),

0 commit comments

Comments
 (0)