commit b25d7fc14491b6f1430b004605cef6d4f6b03995
author Samuel Lee <samuelslee@google.com> Tue Jul 12 10:19:56 2022 -0700
committer TensorFlower Gardener <gardener@tensorflow.org> Tue Jul 12 10:32:47 2022 -0700
tree 39df2ad6716463efbce1f07ee64e88ade7996432
parent b33cfe716deb4e8f97fadf6c370809d318d97c20

Remove Save Op string from prefix string in DTensor SaveV2Op SPMD

During dtensor save op spmd expansion, we add a unique prefix string consisting of device id and the save op. The save op string was previously for tpu or gpu mesh cases since the host mesh was not necessarily 1:1. The host mesh is 1:1 now so the save op number is always 0, i.e each cpu device is running a single save op.

PiperOrigin-RevId: 460494452

tensorflow/dtensor/cc/save_restore_util.cc

diff --git a/tensorflow/dtensor/cc/save_restore_util.cc b/tensorflow/dtensor/cc/save_restore_util.cc
index 0a97c55..2eede56 100644
--- a/tensorflow/dtensor/cc/save_restore_util.cc
+++ b/tensorflow/dtensor/cc/save_restore_util.cc
@@ -144,8 +144,7 @@
         shape_and_slice_specs.push_back({});
         // Generate new prefix based on device_id and save op index, only when
         // we need a new save_op.
-        new_prefixes.push_back(absl::StrCat(prefix, "_device_", device_id,
-                                            "_save_op_", save_op_index));
+        new_prefixes.push_back(absl::StrCat(prefix, "_device_", device_id));
       }
       tensor_indices[save_op_index].push_back(tensor_index);
       shape_and_slice_specs[save_op_index].push_back(specs[save_op_index]);