From b915d3103a38d2f74b0a3aad172d7bfb9ae68e89 Mon Sep 17 00:00:00 2001
From: Hyeontaek Lim <hyeontaek@google.com>
Date: Mon, 20 Oct 2025 14:30:48 -0700
Subject: [PATCH] [NanoRt] NanoRt IFRT now returns a `nullptr` if it knows that
 the Array layout represents a default layout. The user code previously has
 been migrated to handle this new behavior gracefully, obtaining a concrete
 default layout as before.

Future note: At some point, `NanoArray` would need to distinguish between a default layout vs. a concrete layout that is equal to the default layout. If the latter is used, `NanoArray::pjrt_layout()` is expected to return the concrete layout. This is not required by IFRT API semantics yet, but it will be enforced later in the future.

PiperOrigin-RevId: 821808592
---
 third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc      | 6 ++----
 third_party/xla/xla/backends/cpu/nanort/ifrt_client_test.cc | 2 ++
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc b/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc
index e816203685c..dd611d37afc 100644
--- a/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc
+++ b/third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc
@@ -412,9 +412,7 @@ class NanoArray final : public NanoValue<NanoArray, ifrt::Array> {
 
   absl::StatusOr<std::shared_ptr<const PjRtLayout>> pjrt_layout()
       const override {
-    TF_RETURN_IF_ERROR(ValidateNotDeleted());
-    return std::make_shared<PjRtLayout>(
-        LayoutUtil::MakeDescendingLayout(shape().dims().size()));
+    return nullptr;
   }
 
   absl::StatusOr<std::vector<ifrt::ArrayRef>> DisassembleIntoSingleDeviceArrays(
@@ -646,7 +644,7 @@ class ShardedNanoArray final : public NanoValue<ShardedNanoArray, ifrt::Array> {
 
   absl::StatusOr<std::shared_ptr<const PjRtLayout>> pjrt_layout()
       const override {
-    return std::make_shared<PjRtLayout>(Layout(shape().dims()));
+    return nullptr;
   }
 
   absl::StatusOr<std::vector<ifrt::ArrayRef>> DisassembleIntoSingleDeviceArrays(
diff --git a/third_party/xla/xla/backends/cpu/nanort/ifrt_client_test.cc b/third_party/xla/xla/backends/cpu/nanort/ifrt_client_test.cc
index 2d22ff7086e..26d716fe296 100644
--- a/third_party/xla/xla/backends/cpu/nanort/ifrt_client_test.cc
+++ b/third_party/xla/xla/backends/cpu/nanort/ifrt_client_test.cc
@@ -319,6 +319,8 @@ int main(int argc, char** argv) {
       "ArrayImplTest.MakeArrayFromHostBufferAndCopyToHostBufferWithString:"
       "ArrayImplTest."
       "MakeArraysFromHostBufferShardsAndCopyToHostBufferWithString:"
+      // Custom layouts are not supported in NanoIfrtClient.
+      "ArrayImplTest.MakeArraysFromHostBufferShardsWithLayout:"
       // `MakeErrorArrays` is not supported in NanoIfrtClient.
       "ArrayImplTest.MakeErrorArrays:"
       "ArrayImplTest.CopyPoisonedArray:"