Use rank in data propagation (onnx#6557)

### Description

Redo of onnx#6553 (due to DCO/signoff issues)

Extend the implementation of partial data propagation to make use of the
rank of a tensor, if known.

### Motivation and Context

This enables more precise shape inference, especially in the presence of
dynamic shapes involving some unknown dimensions. Generated models often
use computed shapes that mix statically known and statically unknown
dimensions (for example, in an expression such as `Concat (BatchSize,
SequenceLength, OtherKnownDimensions)`). Even if `BatchSize` and
`SequenceLength` are unknown, the other dimensions may be known, and it
useful to track them to enable certain optimizations.

Signed-off-by: Ganesan Ramalingam <[email protected]>

onnx/shape_inference/implementation.h

@@ -413,6 +413,27 @@ struct DataPropagationContextImpl : public DataPropagationContext {
         return &(result.first->second);
       }
     }
+
+    // If X has a known rank N, then X's value can be represented as an array
+    // of N unknown values (represented as a TensorShapeProto).
+    const TypeProto* type = getInputType(index);
+    if ((type != nullptr) && (type->has_tensor_type())) {
+      auto& tensor_type = type->tensor_type();
+      if (tensor_type.has_shape()) {
+        const TensorShapeProto& shape = tensor_type.shape();
+        if ((shape.dim_size() == 1) && (shape.dim(0).has_dim_value())) {
+          TensorShapeProto tsp;
+          int64_t dim_value = shape.dim(0).dim_value();
+          for (int64_t i = 0; i < dim_value; ++i) {
+            tsp.add_dim();
+          }
+          auto result = generatedShapeData_.insert({input_name, std::move(tsp)});
+          if (result.second) {
+            return &(result.first->second);
+          }
+        }
+      }
+    }
     return nullptr;
   }
 

onnx/test/cpp/data_propagation_test.cc

@@ -298,6 +298,23 @@ agraph (int32[1,2] x, int32[3,4] y) => (int32[4] w)
   EXPECT_TRUE(CompareShape(propagated_tsp, expected_tsp));
 }
 
+TEST(DataPropagationImplTest, DynamicConcatTest) {
+  const char* code = R"ONNX(
+agraph (float[32, 1024] x, int64[2] dynamic_shape) => (int64[4] z)
+{
+    xs = Shape(x)   # [32, 1024]
+    z = Concat<axis = 0>(xs, dynamic_shape)  # [32, 1024, ?, ?]
+}
+)ONNX";
+  TensorShapeProto expected_tsp;
+  expected_tsp.mutable_dim()->Add()->set_dim_value(32);
+  expected_tsp.mutable_dim()->Add()->set_dim_value(1024);
+  expected_tsp.mutable_dim()->Add();
+  expected_tsp.mutable_dim()->Add();
+  const auto propagated_tsp = RunDataPropagation(code);
+  EXPECT_TRUE(CompareShape(propagated_tsp, expected_tsp));
+}
+
 TEST(DataPropagationImplTest, GatherTest) {
   const char* code = R"ONNX(
 agraph (int32[1,2,3,4,5,6] x) => (int32[3] w)