When xla::Array::Reshape is called with a dimension span that aliases the array's internal buffer, reallocating the buffer would invalidate the span before the copy was performed. This CL addresses the issue by using a temporary buffer for reallocation and using std::memmove for in-place updates. A regression test 'ReshapeWithLowerRank' is added to verify the fix.

PiperOrigin-RevId: 900935128

Author: Google-ML-Automation

xla/array.h

@@ -542,10 +542,14 @@ class Array {
                         std::multiplies<int64_t>());
     CHECK_EQ(new_num_elements, num_elements());
     if (sizes_.size != new_dimensions.size()) {
-      sizes_ = OwnedBuffer<int64_t>(new_dimensions.size());
+      OwnedBuffer<int64_t> new_sizes(new_dimensions.size());
+      std::memcpy(new_sizes.data.get(), new_dimensions.data(),
+                  new_dimensions.size() * sizeof(int64_t));
+      sizes_ = std::move(new_sizes);
+    } else {
+      std::memmove(sizes_.data.get(), new_dimensions.data(),
+                   new_dimensions.size() * sizeof(int64_t));
     }
-    std::memcpy(sizes_.data.get(), new_dimensions.data(),
-                new_dimensions.size() * sizeof(int64_t));
   }
 
   // Performs a permutation of dimensions.

xla/array_test.cc

@@ -466,5 +466,18 @@ TEST(ArrayTest, UpdateSlice) {
   EXPECT_EQ(expected, arr.ToString());
 }
 
+TEST(ArrayTest, ReshapeWithLowerRank) {
+  Array<int64_t> arr({1, 1, 24});
+  EXPECT_EQ(arr.num_dimensions(), 3);
+  // Reshape to 1D using a sub-span of its own dimensions.
+  // The dropped dimensions are all 1, so the number of elements stays the same.
+  // This triggers reallocation of sizes_ (3 -> 1) while new_dimensions aliases
+  // it.
+  arr.Reshape(arr.dimensions().subspan(2, 1));
+  EXPECT_EQ(arr.num_dimensions(), 1);
+  EXPECT_EQ(arr.dim(0), 24);
+  EXPECT_EQ(arr.num_elements(), 24);
+}
+
 }  // namespace
 }  // namespace xla