From ee240a70ececb723dd376124ea9ca2e4c01aecee Mon Sep 17 00:00:00 2001
From: PragmaTwice <twice@apache.org>
Date: Mon, 16 Dec 2024 16:38:09 +0800
Subject: [PATCH] [PJRT] Fix stablehlo attribute parameters for buffer
 transpose and broadcast

Signed-off-by: PragmaTwice <twice@apache.org>
---
 integrations/pjrt/src/iree_pjrt/common/api_impl.cc | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/integrations/pjrt/src/iree_pjrt/common/api_impl.cc b/integrations/pjrt/src/iree_pjrt/common/api_impl.cc
index 6f2d787ea1d2..67d3baac83e7 100644
--- a/integrations/pjrt/src/iree_pjrt/common/api_impl.cc
+++ b/integrations/pjrt/src/iree_pjrt/common/api_impl.cc
@@ -955,17 +955,17 @@ iree_status_t DeviceInstance::TransposeBroadcastDeviceBuffer(
 
   auto arrayBuilder = [](const int64_t* vals, int64_t sz) {
     std::stringstream ss;
-    ss << " {permutation = dense<[" << vals[0];
+    ss << " {permutation = array<i64: " << vals[0];
     for (int i = 1; i < sz; ++i) ss << ", " << vals[i];
-    ss << "]> : tensor<" << sz << "xi64>}";
+    ss << ">}";
     return ss.str();
   };
 
   auto broadcastBuilder = [](int64_t sz) {
     std::stringstream ss;
-    ss << "{broadcast_dimensions = dense<[0";
+    ss << "{broadcast_dimensions = array<i64: 0";
     for (int i = 1; i < sz; ++i) ss << ", " << i;
-    ss << "]> : tensor<" << sz << "xi64>}";
+    ss << ">}";
     return ss.str();
   };