OSSForks/tensorflow
1
0
Fork 0
mirror of https://github.com/zebrajr/tensorflow.git synced 2025-12-06 12:20:11 +01:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

[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.

Browse Source 
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
This commit is contained in:
Hyeontaek Lim 2025-10-20 14:30:48 -07:00 committed by TensorFlower Gardener
parent dc12ec4556
commit b915d3103a
2 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
third_party/xla/xla/backends/cpu/nanort/ifrt_client.cc vendored
View File

@ -412,9 +412,7 @@ class NanoArray final : public NanoValue<NanoArray, ifrt::Array> {
absl::StatusOr<std::shared_ptr<const PjRtLayout>> pjrt_layout() absl::StatusOr<std::shared_ptr<const PjRtLayout>> pjrt_layout()
const override { const override {
TF_RETURN_IF_ERROR(ValidateNotDeleted()); return nullptr;
return std::make_shared<PjRtLayout>(
LayoutUtil::MakeDescendingLayout(shape().dims().size()));
} }
absl::StatusOr<std::vector<ifrt::ArrayRef>> DisassembleIntoSingleDeviceArrays( 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() absl::StatusOr<std::shared_ptr<const PjRtLayout>> pjrt_layout()
const override { const override {
return std::make_shared<PjRtLayout>(Layout(shape().dims())); return nullptr;
} }
absl::StatusOr<std::vector<ifrt::ArrayRef>> DisassembleIntoSingleDeviceArrays( absl::StatusOr<std::vector<ifrt::ArrayRef>> DisassembleIntoSingleDeviceArrays(

2
third_party/xla/xla/backends/cpu/nanort/ifrt_client_test.cc vendored
View File

@ -319,6 +319,8 @@ int main(int argc, char** argv) {
"ArrayImplTest.MakeArrayFromHostBufferAndCopyToHostBufferWithString:" "ArrayImplTest.MakeArrayFromHostBufferAndCopyToHostBufferWithString:"
"ArrayImplTest." "ArrayImplTest."
"MakeArraysFromHostBufferShardsAndCopyToHostBufferWithString:" "MakeArraysFromHostBufferShardsAndCopyToHostBufferWithString:"
// Custom layouts are not supported in NanoIfrtClient.
"ArrayImplTest.MakeArraysFromHostBufferShardsWithLayout:"
// `MakeErrorArrays` is not supported in NanoIfrtClient. // `MakeErrorArrays` is not supported in NanoIfrtClient.
"ArrayImplTest.MakeErrorArrays:" "ArrayImplTest.MakeErrorArrays:"
"ArrayImplTest.CopyPoisonedArray:" "ArrayImplTest.CopyPoisonedArray:"