From dfc574d5aa4b20382419ef456bf409ca0f13266a Mon Sep 17 00:00:00 2001
From: Changho Hwang <changhohwang@microsoft.com>
Date: Mon, 30 Sep 2024 09:11:14 +0000
Subject: [PATCH] minor fix

---
 ark/api/planner.cpp | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/ark/api/planner.cpp b/ark/api/planner.cpp
index 24036b8f..c2c98b21 100644
--- a/ark/api/planner.cpp
+++ b/ark/api/planner.cpp
@@ -196,13 +196,25 @@ std::string Planner::Impl::plan(bool pretty) const {
 
         auto &result_tensors = op->result_tensors();
         if (!result_tensors.empty() && config.contains("Tile")) {
+            const std::vector<DimType> tile_vec = config["Tile"];
+            std::vector<DimType> trim_leading_ones;
+            for (size_t i = 0; i < tile_vec.size(); i++) {
+                if (tile_vec[i] != 1) {
+                    trim_leading_ones = std::vector<DimType>(
+                        tile_vec.begin() + i, tile_vec.end());
+                    break;
+                }
+            }
+            if (trim_leading_ones.empty()) {
+                trim_leading_ones.push_back(1);
+            }
+            Dims tile(trim_leading_ones);
+
             std::stringstream ss;
-            ss << "Result shape is not divided by tile. Op: "
-               << op->serialize().dump();
+            ss << "Result shape is not divided by tile "
+               << tile << ". Op: " << op->serialize().dump();
             auto not_divided_error = ss.str();
 
-            const std::vector<DimType> tile_vec = config["Tile"];
-            auto tile = Dims(tile_vec);
             auto &result_shape = result_tensors[0]->padded_shape();
             if (result_shape.ndims() < tile.ndims()) {
                 ERR(PlanError, not_divided_error);