[Mlir-commits] [mlir] [mlir][tosa] Fix ability to expand ranks with dynamic shape support (PR #128037)

[llvmlistbot at llvm.org]

================
@@ -77,22 +77,21 @@ computeReshapeOutput(ArrayRef<int64_t> higherRankShape,
   // Initialize new shapes with [1] * higherRank.
   int64_t higherRank = higherRankShape.size();
   int64_t lowerRank = lowerRankShape.size();
-
   reshapeOutputShape.assign(higherRank, 1);
 
   int64_t higherRankDim;
   int64_t lowerRankDim;
+  const int64_t rankDiff = higherRank - lowerRank;
 
-  for (int64_t i = higherRank - 1, j = lowerRank - 1; i >= 0 && j >= 0;
-       i--, j--) {
-    higherRankDim = higherRankShape[i];
-    lowerRankDim = lowerRankShape[j];
+  for (int64_t i = lowerRank - 1; i >= 0; i--) {
+    higherRankDim = higherRankShape[i + rankDiff];
+    lowerRankDim = lowerRankShape[i];
 
-    if (lowerRankDim == 1 && higherRankDim > 1)
-      reshapeOutputShape[i] = 1;
-    else if ((lowerRankDim > 1 && higherRankDim == 1) ||
+    if (lowerRankDim == 1 && higherRankDim != 1)
----------------
Jerry-Ge wrote:

Yes, this can work. Updated the code in the latest commit. 

https://github.com/llvm/llvm-project/pull/128037