Fix an overflow issue in TransposePlan

PiperOrigin-RevId: 715162849

Author: junwhanahn

xla/pjrt/BUILD

@@ -768,6 +768,7 @@ xla_cc_test(
         "//xla:shape_util",
         "//xla:util",
         "//xla/hlo/testlib:test",
+        "//xla/tsl/lib/core:status_test_util",
         "//xla/tsl/protobuf:error_codes_proto_impl_cc",
         "@com_google_absl//absl/container:inlined_vector",
         "@com_google_absl//absl/numeric:int128",

xla/pjrt/transpose.cc

@@ -900,7 +900,7 @@ void TransposePlan::BuildPlanNodes(
 
 absl::StatusOr<std::unique_ptr<TransposePlan>> TransposePlan::Create(
     const Options& o) {
-  auto is_negative = [](int d) { return d < 0; };
+  auto is_negative = [](int64_t d) { return d < 0; };
   if (absl::c_find_if(o.dims, is_negative) != o.dims.end()) {
     return InvalidArgument("dims must be non-negative, got %s",
                            absl::StrJoin(o.dims, ","));

xla/pjrt/transpose_test.cc

@@ -33,6 +33,7 @@ limitations under the License.
 #include "xla/hlo/testlib/test.h"
 #include "xla/permutation_util.h"
 #include "xla/shape_util.h"
+#include "xla/tsl/lib/core/status_test_util.h"
 #include "xla/tsl/protobuf/error_codes.pb.h"
 #include "xla/util.h"
 #include "tsl/platform/statusor.h"
@@ -140,6 +141,20 @@ TEST(TransposeTest, InvalidTilings) {
           "Only one of the input and output may have a non-trivial tiling"));
 }
 
+TEST(TransposeTest, LargeDimensions) {
+  std::vector<int64_t> dims = {3ll << 30};
+  std::vector<int64_t> permutation = {0};
+
+  TransposePlan::Options options;
+  options.elem_size_in_bytes = 8;
+  options.dims = dims;
+  options.permutation = permutation;
+  options.input_layout = TransposePlan::Tiling{};
+  options.output_tiling = TransposePlan::Tiling{};
+  options.transformation = TransposePlan::Transformation::kNone;
+  TF_EXPECT_OK(TransposePlan::Create(options).status());
+}
+
 // Computes the size in elements of a tiled array.
 int64_t SizeOfTiledArray(absl::Span<int64_t const> shape,
                          absl::Span<int64_t const> tiling) {