pjrt_array.cc

/* Copyright 2022 The OpenXLA Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "xla/python/pjrt_ifrt/pjrt_array.h"

#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <variant>
#include <vector>

#include "absl/log/check.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/strings/str_join.h"
#include "absl/types/span.h"
#include "xla/literal.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/pjrt/pjrt_layout.h"
#include "xla/pjrt/utils.h"
#include "xla/python/ifrt/array.h"
#include "xla/python/ifrt/device.h"
#include "xla/python/ifrt/dtype.h"
#include "xla/python/ifrt/future.h"
#include "xla/python/ifrt/memory.h"
#include "xla/python/ifrt/shape.h"
#include "xla/python/ifrt/sharding.h"
#include "xla/python/pjrt_ifrt/pjrt_client.h"
#include "xla/shape_util.h"
#include "xla/status.h"
#include "xla/util.h"
#include "xla/xla_data.pb.h"
#include "tsl/concurrency/ref_count.h"
#include "tsl/platform/errors.h"
#include "tsl/platform/statusor.h"

namespace xla {
namespace ifrt {

namespace {

// Validates the sharding and PjRtBuffers have consistent device and memory
// kind.
Status ValidateArrayCreationInput(std::shared_ptr<const Sharding> sharding,
                                  const PjRtArray::PjRtBuffers& pjrt_buffers) {
  if (pjrt_buffers.empty()) {
    return InvalidArgument("pjrt_buffers must be non-empty");
  }
  if (sharding->devices().size() != pjrt_buffers.size()) {
    return InvalidArgument("device and buffer counts mismatch: %d vs. %d",
                           sharding->devices().size(), pjrt_buffers.size());
  }

  // Canonicalize memory kind in case it hasn't been done before.
  MemoryKind canonicalized_sharding_memory_kind = CanonicalizeMemoryKind(
      sharding->memory_kind(), sharding->devices().front());
  for (int i = 0; i < sharding->devices().size(); ++i) {
    if (pjrt_buffers[i]->device() != sharding->devices()[i]) {
      return InvalidArgument(
          "PjRtBuffer's memory space is addressed by device %s vs sharding is "
          "on device %s",
          pjrt_buffers[i]->device()->DebugString(),
          sharding->devices()[i]->DebugString());
    }
    MemoryKind buffer_memory_kind =
        MakeMemoryKindFromPjRtBuffer(pjrt_buffers[i].get());
    if (canonicalized_sharding_memory_kind != buffer_memory_kind) {
      return InvalidArgument(
          "PjRtBuffer's memory kind does not match sharding's memory kind. Got "
          "PjRtBuffer's memory kind: %s vs shardings's memory kind: %s",
          buffer_memory_kind.DebugString(),
          canonicalized_sharding_memory_kind.DebugString());
    }
  }
  return OkStatus();
}

// Validates the PjRtBuffers have consistent memory kind and returns the memory
// kind.
absl::StatusOr<MemoryKind> GetMemoryKindFromPjRtBuffers(
    const PjRtArray::PjRtBuffers& pjrt_buffers) {
  const auto first_memory_kind =
      MakeMemoryKindFromPjRtBuffer(pjrt_buffers.front().get());
  const MemoryKind canonical_first_memory_kind =
      CanonicalizeMemoryKind(first_memory_kind, pjrt_buffers.front()->device());
  for (const auto& pjrt_buffer : pjrt_buffers) {
    if (auto memory_kind = MakeMemoryKindFromPjRtBuffer(pjrt_buffer.get());
        canonical_first_memory_kind !=
        CanonicalizeMemoryKind(memory_kind, pjrt_buffer->device())) {
      return InvalidArgument(
          "Memory kind mismatch between PjRtBuffers. Got one buffer with "
          "memory kind: %s and another with memory_kind: %s",
          first_memory_kind.DebugString(), memory_kind.DebugString());
    }
  }
  return first_memory_kind;
}

}  // namespace

char PjRtCompatibleArray::ID = 0;
char PjRtArray::ID = 0;

absl::StatusOr<xla::PrimitiveType> ToPrimitiveType(DType dtype) {
  switch (dtype.kind()) {
#define CASE(DT, PT)                                                      \
  case DT:                                                                \
    static_assert(PT ==                                                   \
                  static_cast<xla::PrimitiveType>(static_cast<int>(DT))); \
    return PT
    CASE(DType::kInvalid, xla::PrimitiveType::PRIMITIVE_TYPE_INVALID);
    CASE(DType::kPred, xla::PrimitiveType::PRED);
    CASE(DType::kS4, xla::PrimitiveType::S4);
    CASE(DType::kS8, xla::PrimitiveType::S8);
    CASE(DType::kS16, xla::PrimitiveType::S16);
    CASE(DType::kS32, xla::PrimitiveType::S32);
    CASE(DType::kS64, xla::PrimitiveType::S64);
    CASE(DType::kU4, xla::PrimitiveType::U4);
    CASE(DType::kU8, xla::PrimitiveType::U8);
    CASE(DType::kU16, xla::PrimitiveType::U16);
    CASE(DType::kU32, xla::PrimitiveType::U32);
    CASE(DType::kU64, xla::PrimitiveType::U64);
    CASE(DType::kF8E4M3FN, xla::PrimitiveType::F8E4M3FN);
    CASE(DType::kF8E4M3B11FNUZ, xla::PrimitiveType::F8E4M3B11FNUZ);
    CASE(DType::kF8E4M3FNUZ, xla::PrimitiveType::F8E4M3FNUZ);
    CASE(DType::kF8E5M2, xla::PrimitiveType::F8E5M2);
    CASE(DType::kF8E5M2FNUZ, xla::PrimitiveType::F8E5M2FNUZ);
    CASE(DType::kF16, xla::PrimitiveType::F16);
    CASE(DType::kF32, xla::PrimitiveType::F32);
    CASE(DType::kBF16, xla::PrimitiveType::BF16);
    CASE(DType::kF64, xla::PrimitiveType::F64);
    CASE(DType::kC64, xla::PrimitiveType::C64);
    CASE(DType::kC128, xla::PrimitiveType::C128);
    CASE(DType::kToken, xla::PrimitiveType::TOKEN);
#undef CASE
    case DType::kString:
      return InvalidArgument("Not supported as XLA PrimitiveType: %d",
                             static_cast<int>(dtype.kind()));
  }
  return InvalidArgument("Invalid DType: %d", static_cast<int>(dtype.kind()));
}

absl::StatusOr<DType> ToDType(xla::PrimitiveType primitive_type) {
  switch (primitive_type) {
    case xla::PrimitiveType::PRIMITIVE_TYPE_INVALID:
    case xla::PrimitiveType::PRED:
    case xla::PrimitiveType::S4:
    case xla::PrimitiveType::S8:
    case xla::PrimitiveType::S16:
    case xla::PrimitiveType::S32:
    case xla::PrimitiveType::S64:
    case xla::PrimitiveType::U4:
    case xla::PrimitiveType::U8:
    case xla::PrimitiveType::U16:
    case xla::PrimitiveType::U32:
    case xla::PrimitiveType::U64:
    case xla::PrimitiveType::F8E4M3FN:
    case xla::PrimitiveType::F8E4M3B11FNUZ:
    case xla::PrimitiveType::F8E4M3FNUZ:
    case xla::PrimitiveType::F8E5M2:
    case xla::PrimitiveType::F8E5M2FNUZ:
    case xla::PrimitiveType::F16:
    case xla::PrimitiveType::F32:
    case xla::PrimitiveType::BF16:
    case xla::PrimitiveType::F64:
    case xla::PrimitiveType::C64:
    case xla::PrimitiveType::C128:
    case xla::PrimitiveType::TOKEN:
      return DType(static_cast<DType::Kind>(static_cast<int>(primitive_type)));
    default:
      return InvalidArgument("Invalid XLA PrimitiveType: %d",
                             static_cast<int>(primitive_type));
  }
}

MemoryKind MakeMemoryKindFromPjRtBuffer(PjRtBuffer* pjrt_buffer) {
  if (pjrt_buffer->memory_space() == nullptr) {
    return MemoryKind();
  }
  return MemoryKind(pjrt_buffer->memory_space()->kind());
}

absl::StatusOr<tsl::RCReference<PjRtArray>> PjRtArray::Create(
    PjRtCompatibleClient* client, DType dtype, Shape shape,
    std::shared_ptr<const Sharding> sharding, PjRtBuffers pjrt_buffers) {
  TF_RETURN_IF_ERROR(ValidateArrayCreationInput(sharding, pjrt_buffers));
  return tsl::MakeRef<PjRtArray>(client, dtype, std::move(shape),
                                 std::move(sharding), std::move(pjrt_buffers));
}

absl::StatusOr<tsl::RCReference<PjRtArray>> PjRtArray::Create(
    PjRtCompatibleClient* client, DType dtype, DynamicShape dynamic_shape,
    std::shared_ptr<const Sharding> sharding, PjRtBuffers pjrt_buffers) {
  TF_RETURN_IF_ERROR(ValidateArrayCreationInput(sharding, pjrt_buffers));
  return tsl::MakeRef<PjRtArray>(client, dtype, std::move(dynamic_shape),
                                 std::move(sharding), std::move(pjrt_buffers));
}

absl::StatusOr<tsl::RCReference<PjRtArray>> PjRtArray::Create(
    PjRtCompatibleClient* client, std::shared_ptr<PjRtBuffer> pjrt_buffer) {
  TF_ASSIGN_OR_RETURN(auto dtype, ToDType(pjrt_buffer->element_type()));
  Shape shape(pjrt_buffer->dimensions());
  auto sharding = SingleDeviceSharding::Create(
      pjrt_buffer->device(), MakeMemoryKindFromPjRtBuffer(pjrt_buffer.get()));
  return tsl::MakeRef<PjRtArray>(client, dtype, std::move(shape),
                                 std::move(sharding),
                                 PjRtBuffers({std::move(pjrt_buffer)}));
}

absl::StatusOr<tsl::RCReference<Array>> PjRtArray::FullyReplicatedShard(
    ArrayCopySemantics semantics) {
  return PjRtArray::Create(client(), GetPjRtBuffer(semantics, 0));
}

std::shared_ptr<PjRtBuffer> PjRtArray::GetPjRtBuffer(
    ArrayCopySemantics semantics, int index) const {
  switch (semantics) {
    case ArrayCopySemantics::kAlwaysCopy:
      // TODO(hyeontaek): kAlwaysCopy should clone the buffer, but the PjRt
      // API does not have efficient buffer cloning on the same device.
      return pjrt_buffers_[index];
    case ArrayCopySemantics::kReuseInput:
      return pjrt_buffers_[index];
    case ArrayCopySemantics::kDonateInput:
      // TODO(hyeontaek): We may try std::move(pjrt_buffers_[i]), but this
      // would be unsafe if there is a subsequent access to the buffer.
      return pjrt_buffers_[index];
  }
}

absl::StatusOr<tsl::RCReference<PjRtArray>> PjRtArray::Create(
    PjRtCompatibleClient* client, Shape shape, PjRtBuffers pjrt_buffers) {
  TF_ASSIGN_OR_RETURN(auto dtype,
                      xla::ifrt::ToDType(pjrt_buffers.front()->element_type()));
  TF_ASSIGN_OR_RETURN(MemoryKind memory_kind,
                      GetMemoryKindFromPjRtBuffers(pjrt_buffers));

  DeviceList::Devices devices;
  devices.reserve(pjrt_buffers.size());
  std::vector<Shape> shapes;
  shapes.reserve(pjrt_buffers.size());

  for (const auto& pjrt_buffer : pjrt_buffers) {
    devices.push_back(pjrt_buffer->device());
    shapes.push_back(Shape(pjrt_buffer->dimensions()));
  }
  auto sharding = ifrt::ConcreteSharding::Create(DeviceList(std::move(devices)),
                                                 memory_kind,
                                                 /*shape=*/shape,
                                                 /*shard_shapes=*/shapes);
  return PjRtArray::Create(client, dtype, std::move(shape), std::move(sharding),
                           std::move(pjrt_buffers));
}

absl::StatusOr<tsl::RCReference<PjRtArray>> PjRtArray::Create(
    PjRtCompatibleClient* client, DynamicShape dynamic_shape,
    PjRtBuffers pjrt_buffers) {
  TF_ASSIGN_OR_RETURN(auto dtype,
                      xla::ifrt::ToDType(pjrt_buffers.front()->element_type()));
  TF_ASSIGN_OR_RETURN(auto memory_kind,
                      GetMemoryKindFromPjRtBuffers(pjrt_buffers));

  DeviceList::Devices devices;
  devices.reserve(pjrt_buffers.size());
  std::vector<DynamicShape> dynamic_shapes;
  dynamic_shapes.reserve(pjrt_buffers.size());

  for (const auto& pjrt_buffer : pjrt_buffers) {
    devices.push_back(pjrt_buffer->device());
    TF_ASSIGN_OR_RETURN(
        DynamicShape dynamic_shape,
        // Extracts dynamic shape info from the buffers.
        DynamicShape::Create(
            Shape(pjrt_buffer->dimensions()),
            BoundedDynamicShapeTag(pjrt_buffer->is_dynamic_dimension())));
    dynamic_shapes.push_back(std::move(dynamic_shape));
  }
  auto sharding = ifrt::ConcreteSharding::Create(
      DeviceList(std::move(devices)), memory_kind,
      /*dynamic_shape=*/dynamic_shape,
      /*shard_dynamic_shapes=*/dynamic_shapes);
  return PjRtArray::Create(client, dtype, std::move(dynamic_shape),
                           std::move(sharding), std::move(pjrt_buffers));
}

PjRtArray::PjRtArray(PjRtCompatibleClient* client, DType dtype, Shape shape,
                     std::shared_ptr<const Sharding> sharding,
                     PjRtBuffers pjrt_buffers)
    : client_(client),
      dtype_(dtype),
      shape_(std::move(shape)),
      sharding_(std::move(sharding)),
      pjrt_buffers_(std::move(pjrt_buffers)) {}

PjRtArray::PjRtArray(PjRtCompatibleClient* client, DType dtype,
                     DynamicShape dynamic_shape,
                     std::shared_ptr<const Sharding> sharding,
                     PjRtBuffers pjrt_buffers)
    : client_(client),
      dtype_(dtype),
      shape_(std::move(dynamic_shape)),
      sharding_(std::move(sharding)),
      pjrt_buffers_(std::move(pjrt_buffers)) {}

absl::StatusOr<std::vector<tsl::RCReference<Array>>>
PjRtArray::DisassembleIntoSingleDeviceArrays(ArrayCopySemantics semantics) {
  DCHECK(this);
  std::vector<tsl::RCReference<Array>> result;
  result.reserve(sharding_->devices().size());
  TF_RETURN_IF_ERROR(std::visit(
      [&](const auto& this_shape) {
        TF_ASSIGN_OR_RETURN(auto shape_and_shardings,
                            sharding_->Disassemble(this_shape));
        for (int i = 0; i < sharding_->devices().size(); ++i) {
          PjRtBuffers buffers;
          buffers.reserve(1);
          buffers.push_back(GetPjRtBuffer(semantics, i));
          TF_ASSIGN_OR_RETURN(
              auto array,
              PjRtArray::Create(client_, dtype_,
                                std::move(shape_and_shardings[i].first),
                                std::move(shape_and_shardings[i].second),
                                std::move(buffers)));
          result.push_back(std::move(array));
        }
        return xla::OkStatus();
      },
      shape_));

  return result;
}

Future<Status> PjRtArray::CopyToHostBuffer(
    void* data, std::optional<absl::Span<const int64_t>> byte_strides,
    ArrayCopySemantics semantics) {
  DCHECK(this);
  if (sharding_->devices().size() != 1) {
    return Future<Status>(
        InvalidArgument("Only single-shard is implemented, but got %d",
                        sharding_->devices().size()));
  }

  auto dtype = ToPrimitiveType(dtype_);
  if (!dtype.ok()) {
    return Future<Status>(std::move(dtype).status());
  }

  PjRtBuffer* pjrt_buffer = pjrt_buffers_.front().get();
  absl::Span<const int64_t> dims;
  absl::StatusOr<std::vector<int64_t>> logical_dims;
  if (!pjrt_buffer->has_dynamic_dimensions()) {
    dims = std::get<Shape>(shape_).dims();
  } else {
    // TODO(b/182461453): This is a blocking call. If we further implemented
    // populating dynamic shape metadata while fetching the literal, we wouldn't
    // need this static approach.
    // TODO(hyeontaek): Clean up this dynamic shape access once we formalize
    // dynamic shape support in IFRT.
    // TODO(b/314805296): Use the new dynamic shape here.
    logical_dims = pjrt_buffer->logical_dimensions();
    if (!logical_dims.ok()) {
      return Future<Status>(std::move(logical_dims).status());
    }
    dims = *logical_dims;
  }

  std::unique_ptr<xla::MutableBorrowingLiteral> literal;
  if (byte_strides.has_value()) {
    auto xla_shape =
        MakeShapeWithTrivialByteStrides(*dtype, dims, *byte_strides);
    if (!xla_shape.ok()) {
      return Future<Status>(std::move(xla_shape).status());
    }
    literal = std::make_unique<xla::MutableBorrowingLiteral>(
        static_cast<char*>(data), *xla_shape);
  } else {
    auto xla_shape = ShapeUtil::MakeShapeWithDescendingLayout(*dtype, dims);
    literal = std::make_unique<xla::MutableBorrowingLiteral>(
        static_cast<char*>(data), xla_shape);
  }
  auto* literal_ptr = literal.get();
  auto promise = Future<Status>::CreatePromise();
  Future<Status> future(promise);
  // TODO(hyeontaek): Handle semantics == kDonateInput.
  pjrt_buffer->ToLiteral(literal_ptr)
      .OnReady([literal = std::move(literal),
                promise = std::move(promise)](Status s) mutable {
        promise.Set(std::move(s));
        literal = nullptr;
      });
  return future;
}

// TODO(yashkatariya): Maybe move this to ifrt::Device?
absl::StatusOr<PjRtMemorySpace*> GetMemorySpaceFromMemoryKind(
    ifrt::Device* device, ifrt::MemoryKind memory_kind) {
  PjRtMemorySpace* memory_space = nullptr;
  for (PjRtMemorySpace* ms : device->memory_spaces()) {
    if (ms->kind() == memory_kind.memory_kind()) {
      memory_space = ms;
      break;
    }
  }
  if (memory_space == nullptr) {
    return InvalidArgument(
        "Invalid memory kind: %s; available memory kinds: %s",
        memory_kind.DebugString(),
        absl::StrJoin(device->memory_spaces(), ", ",
                      [](std::string* out, PjRtMemorySpace* ms) {
                        absl::StrAppend(out, ms->kind());
                      }));
  }
  return memory_space;
}

absl::StatusOr<tsl::RCReference<Array>> PjRtArray::Reshard(
    std::shared_ptr<const Sharding> new_sharding,
    ArrayCopySemantics semantics) {
  DCHECK(this);
  if (new_sharding->devices().size() != sharding_->devices().size()) {
    return InvalidArgument(
        "Resharding to a different number of devices: %d; expected %d",
        new_sharding->devices().size(), sharding_->devices().size());
  }
  // TODO(hyeontaek): We should have an equivalence test for sharding that
  // permits device changes and nothing else.
  PjRtBuffers buffers;
  buffers.reserve(pjrt_buffers_.size());
  CHECK_GT(new_sharding->devices().size(), 0);
  // Canonicalize memory kind in case it hasn't been done before.
  MemoryKind canonicalized_sharding_memory_kind = CanonicalizeMemoryKind(
      new_sharding->memory_kind(), new_sharding->devices().front());
  bool new_sharding_has_memory_kind =
      canonicalized_sharding_memory_kind.memory_kind().has_value();
  for (int i = 0; i < pjrt_buffers_.size(); ++i) {
    bool devices_equal =
        pjrt_buffers_[i]->device() == new_sharding->devices()[i];
    bool memories_supported = pjrt_buffers_[i]->memory_space() != nullptr;
    bool memory_kind_equal =
        new_sharding_has_memory_kind && memories_supported &&
        pjrt_buffers_[i]->memory_space()->kind() ==
            canonicalized_sharding_memory_kind.memory_kind();

    // No need for data transfer.
    if (devices_equal && (!new_sharding_has_memory_kind ||
                          !memories_supported || memory_kind_equal)) {
      switch (semantics) {
        case ArrayCopySemantics::kAlwaysCopy:
          // TODO(hyeontaek): kAlwaysCopy should clone the buffer, but the PjRt
          // API does not have efficient buffer cloning on the same device.
          buffers.push_back(pjrt_buffers_[i]);
          break;
        case ArrayCopySemantics::kReuseInput:
          buffers.push_back(pjrt_buffers_[i]);
          break;
        case ArrayCopySemantics::kDonateInput:
          // TODO(hyeontaek): We may try std::move(pjrt_buffers_[i]), but this
          // would be unsafe if there is a subsequent access to the buffer.
          buffers.push_back(pjrt_buffers_[i]);
          break;
      }
    } else {
      if (new_sharding->devices()[i]->client() == nullptr) {
        return InvalidArgument(
            "The destination device is owned by a non-PjRt-compatible client. "
            "To use this Array on the destination device, the Array must be "
            "first fetched to the host and then sent to the destination "
            "device.");
      }
      // Use `PjRtBuffer::CopyToMemorySpace` instead of
      // `PjRtBuffer::CopyToDevice` when memories are supported. Because the
      // semantics of the latter one is to copy to the default memory space of
      // the device.
      if (new_sharding_has_memory_kind && memories_supported) {
        TF_ASSIGN_OR_RETURN(
            auto memory_space,
            GetMemorySpaceFromMemoryKind(new_sharding->devices()[i],
                                         canonicalized_sharding_memory_kind));
        TF_ASSIGN_OR_RETURN(std::unique_ptr<PjRtBuffer> copied_buffer,
                            pjrt_buffers_[i]->CopyToMemorySpace(memory_space));
        if (semantics == ArrayCopySemantics::kDonateInput) {
          if (!memory_kind_equal) {
            return Unimplemented(
                "Donation across different memory kinds is not implemented.");
          }
          pjrt_buffers_[i] = nullptr;
        }
        buffers.push_back(std::shared_ptr<PjRtBuffer>(copied_buffer.release()));
      } else {
        // Use `PjRtBuffer::CopyToDevice` when memories are not supported.
        TF_ASSIGN_OR_RETURN(
            std::unique_ptr<xla::PjRtBuffer> copied_buffer,
            pjrt_buffers_[i]->CopyToDevice(new_sharding->devices()[i]));
        if (semantics == ArrayCopySemantics::kDonateInput) {
          pjrt_buffers_[i] = nullptr;
        }
        buffers.push_back(std::shared_ptr<PjRtBuffer>(copied_buffer.release()));
      }
    }
  }
  return std::visit(
      [this, &new_sharding, &buffers](const auto& shape) {
        return PjRtArray::Create(client_, dtype_, shape,
                                 std::move(new_sharding), std::move(buffers));
      },
      shape_);
}

Future<Status> PjRtArray::GetReadyFuture() const {
  DCHECK(this);
  if (pjrt_buffers_.size() == 1) {
    return pjrt_buffers_.front()->GetReadyFuture().ToStatusFuture();
  }
  std::vector<Future<Status>> futures;
  futures.reserve(pjrt_buffers_.size());
  for (auto& buf : pjrt_buffers_) {
    futures.push_back(buf->GetReadyFuture().ToStatusFuture());
  }
  return JoinFutures(absl::MakeSpan(futures));
}

Future<Status> PjRtArray::Delete() {
  DCHECK(this);
  for (auto& buffer : pjrt_buffers_) {
    buffer->Delete();
  }
  // TODO(hyeontaek): Return a correct future.
  return Future<Status>(OkStatus());
}

bool PjRtArray::IsDeleted() const {
  DCHECK(this);
  // TODO(hyeontaek): This may be incorrect if PjRtBuffers are shared and a
  // portion of pjrt_buffers_ is deleted or not deleted.
  return pjrt_buffers_.front()->IsDeleted();
}

std::string PjRtArray::DebugString() const {
  DCHECK(this);
  absl::StatusOr<std::unique_ptr<PjRtLayout>> layout_ptr = layout();
  std::string layout_str =
      layout_ptr.ok() ? (*layout_ptr)->ToString() : "<unknown>";

  return absl::StrFormat(
      "PjRtArray(dtype=%s; shape=%s; sharding=%s; layout=%s)",
      dtype_.DebugString(),
      std::visit([](const auto& shape) { return shape.DebugString(); }, shape_),
      sharding_->DebugString(), layout_str);
}

// TODO(b/330198879): populate layout at construction instead of accessing PJRT
// buffer directly for consistency with Pathways.
absl::StatusOr<std::unique_ptr<PjRtLayout>> PjRtArray::layout() const {
  CHECK(!pjrt_buffers_.empty());
  std::unique_ptr<PjRtLayout> layout = pjrt_buffers_[0]->layout();
#ifndef NDEBUG
  for (int i = 1; i < pjrt_buffers_.size(); ++i) {
    std::unique_ptr<PjRtLayout> layout_i = pjrt_buffers_[i]->layout();
    DCHECK(*layout == *layout_i)
        << "PjRtArray has mismatched layouts across shards! "
        << "shard 0: " << layout->ToString() << ", shard " << i << ": "
        << layout_i->ToString();
  }
#endif
  return layout;
}

}  // namespace ifrt
}  // namespace xla