diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.cc b/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.cc
index c92a6dfbd3890f..14f3dd69c0ed37 100644
--- a/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.cc
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.cc
@@ -91,6 +91,7 @@ void RamFileBlockCache::UpdateLRU(const Key& key,
   // in the cache, and our current block is not block size, this likely means
   // we have inconsistent state within the cache. Note: it's possible some
   // incomplete reads may still go undetected.
+  // Read of block->data allowed because block_state==FINISHED here.
   if (block->data.size() < block_size_) {
     Key fmax = std::make_pair(key.first, std::numeric_limits<size_t>::max());
     auto fcmp = block_map_.upper_bound(fmax);
@@ -105,6 +106,8 @@ void RamFileBlockCache::UpdateLRU(const Key& key,
   return TF_SetStatus(status, TF_OK, "");
 }
 
+// Attempt to fetch data with the given key into block.
+// Return with *status TF_OK only if block->state==FINISHED.
 void RamFileBlockCache::MaybeFetch(const Key& key,
                                    const std::shared_ptr<Block>& block,
                                    TF_Status* status) {
@@ -113,10 +116,12 @@ void RamFileBlockCache::MaybeFetch(const Key& key,
     // Perform this action in a cleanup callback to avoid locking mu_ after
     // locking block->mu.
     if (downloaded_block) {
+      // downloaded_block == (block->state==FINISHED), so reads of block->data
+      // are legal here
       absl::MutexLock l(mu_);
       // Do not update state if the block is already to be evicted.
       if (block->timestamp != 0) {
-        // Use capacity() instead of size() to account for all  memory
+        // Use capacity() instead of size() to account for all memory
         // used by the cache.
         cache_size_ += block->data.capacity();
         // Put to beginning of LRA list.
@@ -137,6 +142,7 @@ void RamFileBlockCache::MaybeFetch(const Key& key,
         // TF_FALLTHROUGH_INTENDED
       case FetchState::CREATED:
         block->state = FetchState::FETCHING;
+        // Thread may modify block->data block->state==FETCHING.
         block->mu.unlock();  // Release the lock while making the API call.
         block->data.clear();
         block->data.resize(block_size_, 0);
@@ -202,6 +208,8 @@ int64_t RamFileBlockCache::Read(const std::string& filename, size_t offset,
     }
     MaybeFetch(key, block, status);
     if (TF_GetCode(status) != TF_OK) return -1;
+    // At this point, block->state==FINISHED, since MaybeFetch() yielded TF_OK.
+    // Therefore, it is legal to access block->data.
     UpdateLRU(key, block, status);
     if (TF_GetCode(status) != TF_OK) return -1;
     // Copy the relevant portion of the block into the result buffer.
@@ -282,10 +290,14 @@ void RamFileBlockCache::Prune() {
 
 void RamFileBlockCache::Flush() {
   absl::MutexLock lock(mu_);
-  block_map_.clear();
-  lru_list_.clear();
-  lra_list_.clear();
-  cache_size_ = 0;
+  // This code mirrors that in RemoveFile_Locked(),
+  // but iterates over the entire cache.
+  auto it = block_map_.begin();
+  while (it != block_map_.end()) {
+    auto next = std::next(it);
+    RemoveBlock(it);
+    it = next;
+  }
 }
 
 void RamFileBlockCache::RemoveFile(const std::string& filename) {
@@ -309,7 +321,17 @@ void RamFileBlockCache::RemoveBlock(BlockMap::iterator entry) {
   entry->second->timestamp = 0;
   lru_list_.erase(entry->second->lru_iterator);
   lra_list_.erase(entry->second->lra_iterator);
-  cache_size_ -= entry->second->data.capacity();
+
+  // Adjust the cache_size_ by the size of the block.
+  // RemoveBlock() can be called by Flush() on blocks that are not yet FINISHED.
+  // Only finished blocks are counted in cache_size_, and it would be a race
+  // to read the data member of a block that is not yet FINISHED.
+  entry->second->mu.lock();
+  if (entry->second->state == FetchState::FINISHED) {
+    cache_size_ -= entry->second->data.capacity();
+  }
+  entry->second->mu.unlock();
+
   block_map_.erase(entry);
 }
 
diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h b/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h
index 3e972fa6292995..86558a295fbd64 100644
--- a/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache.h
@@ -180,8 +180,10 @@ class RamFileBlockCache {
   /// Thread safety:
   /// The iterator and timestamp fields should only be accessed while holding
   /// the block-cache-wide mu_ instance variable. The state variable should only
-  /// be accessed while holding the Block's mu lock. The data vector should only
-  /// be accessed after state == FINISHED, and it should never be modified.
+  /// be accessed while holding the Block's mu lock.  Once state==FINISHED,
+  /// the data vector may be read and may not be written; before
+  /// state==FINISHED, the data vector may be accessed only by the thread that
+  /// set state==FETCHING.
   ///
   /// In order to prevent deadlocks, never grab the block-cache-wide mu_ lock
   /// AFTER grabbing any block's mu lock. It is safe to grab mu without locking
@@ -212,13 +214,13 @@ class RamFileBlockCache {
   void Prune() ABSL_LOCKS_EXCLUDED(mu_);
 
   bool BlockNotStale(const std::shared_ptr<Block>& block)
-      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) ABSL_LOCKS_EXCLUDED(block->mu);
 
   /// Look up a Key in the block cache.
   std::shared_ptr<Block> Lookup(const Key& key) ABSL_LOCKS_EXCLUDED(mu_);
 
   void MaybeFetch(const Key& key, const std::shared_ptr<Block>& block,
-                  TF_Status* status) ABSL_LOCKS_EXCLUDED(mu_);
+                  TF_Status* status) ABSL_LOCKS_EXCLUDED(mu_, block->mu);
 
   /// Trim the block cache to make room for another entry.
   void Trim() ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_);
diff --git a/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache_test.cc b/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache_test.cc
index 18438ef0580f72..929fc25a6372bb 100644
--- a/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache_test.cc
+++ b/tensorflow/c/experimental/filesystem/plugins/gcs/ram_file_block_cache_test.cc
@@ -610,5 +610,34 @@ TEST(RamFileBlockCacheTest, Flush) {
   EXPECT_EQ(calls, 2);
 }
 
+// This test case reproduces a bug, now fixed, in which
+// RamFileBlockCache::Flush() used to clear the cache, invalidating the
+// lra_iterator field, without resetting the timestamp fields of the blocks.
+// This was a problem for the "reconcile_state" cleanup callback in
+// MaybeFetch(), which assumes that if the timestamp field is non-zero, the
+// lra_iterator is valid.  Later versions of Flush() should handle this case
+// correctly.
+TEST(RamFileBlockCacheTest, FlushDuringFetch) {
+  // The delays are sized to exceed the delay in RamFileBlockCache::Prune().
+  auto delayed_fetcher = [](const std::string& filename, size_t offset,
+                            size_t n, char* buffer,
+                            TF_Status* status) -> int64_t {
+    memset(buffer, 'x', n);
+    TF_SetStatus(status, TF_OK, "");
+    Env::Default()->SleepForMicroseconds(10 * 1000 * 1000);
+    return n;
+  };
+  tf_gcs_filesystem::RamFileBlockCache cache(16, 32, /*max_staleness=*/20,
+                                             delayed_fetcher);
+  std::vector<char> out;
+  std::unique_ptr<Thread> flush_thread(
+      Env::Default()->StartThread({}, "delayed_flush", [&cache] {
+        Env::Default()->SleepForMicroseconds(5 * 1000 * 1000);
+        cache.Flush();
+      }));
+  TF_EXPECT_OK(ReadCache(&cache, "", 0, 16, &out));
+  Env::Default()->SleepForMicroseconds(5 * 1000 * 1000);
+}
+
 }  // namespace
 }  // namespace tensorflow
diff --git a/tensorflow/lite/CMakeLists.txt b/tensorflow/lite/CMakeLists.txt
index 39599a21cc5681..27e25e83750461 100644
--- a/tensorflow/lite/CMakeLists.txt
+++ b/tensorflow/lite/CMakeLists.txt
@@ -424,14 +424,14 @@ if(TFLITE_ENABLE_GPU)
     enable_language(OBJCXX)
     list(APPEND TFLITE_DELEGATES_METAL_SRCS
       ${TFLITE_SOURCE_DIR}/delegates/gpu/metal_delegate.mm
-      ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/buffer.cc
+      ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/buffer.mm
       ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/buffer_convert.mm
       ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/common.mm
-      ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/compute_task.cc
-      ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context.cc
-      ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/metal_arguments.cc
-      ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/metal_device.cc
-      ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/metal_spatial_tensor.cc
+      ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/compute_task.mm
+      ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context.mm
+      ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/metal_arguments.mm
+      ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/metal_device.mm
+      ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/metal_spatial_tensor.mm
     )
     add_library(metal_delegate STATIC
       ${TFLITE_DELEGATES_METAL_SRCS}
@@ -440,6 +440,7 @@ if(TFLITE_ENABLE_GPU)
       ${CMAKE_BINARY_DIR}/abseil-cpp
       ${CMAKE_BINARY_DIR}/flatbuffers/include
       PRIVATE ${TENSORFLOW_SOURCE_DIR}
+      PRIVATE ${PROJECT_BINARY_DIR}
     )
     #
     # generate flatbuffers header for inference_context
@@ -449,54 +450,22 @@ if(TFLITE_ENABLE_GPU)
     else()
       set(FLATC flatc)
     endif()
+    set(METAL_GENERATED_DIR ${PROJECT_BINARY_DIR}/tensorflow/lite/delegates/gpu/metal)
     add_custom_command(
-        OUTPUT ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context_generated.h
+        OUTPUT ${METAL_GENERATED_DIR}/inference_context_generated.h
+        COMMAND ${CMAKE_COMMAND} -E make_directory ${METAL_GENERATED_DIR}
         COMMAND ${FLATC} --scoped-enums
         -I ${TENSORFLOW_SOURCE_DIR}
-        -o ${TFLITE_SOURCE_DIR}/delegates/gpu/metal
+        -o ${METAL_GENERATED_DIR}
         -c ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context.fbs
+        DEPENDS ${FLATC_TARGET}
+                ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context.fbs
     )
     add_custom_target(
         inference_context_cc_fbs
-        DEPENDS ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/inference_context_generated.h
+        DEPENDS ${METAL_GENERATED_DIR}/inference_context_generated.h
     )
     add_dependencies(metal_delegate inference_context_cc_fbs)
-    #
-    # supplementary libraries for libmetal_delegate
-    #
-    list(APPEND CC_SRCS
-        buffer
-        compute_task
-        inference_context
-        metal_arguments
-        metal_device
-        metal_spatial_tensor
-    )
-   SET(METAL_DELEGATE_PATH ${TFLITE_SOURCE_DIR}/delegates/gpu/metal/)
-
-   foreach(lib_name ${CC_SRCS})
-     set_source_files_properties(${METAL_DELEGATE_PATH}${lib_name}.cc  PROPERTIES LANGUAGE OBJCXX)
-     add_library("${lib_name}" STATIC ${METAL_DELEGATE_PATH}${lib_name}.cc)
-     target_include_directories("${lib_name}" PUBLIC
-       ${CMAKE_BINARY_DIR}/abseil-cpp
-       ${CMAKE_BINARY_DIR}/flatbuffers/include
-     )
-     set_target_properties(${lib_name} PROPERTIES LINKER_LANGUAGE OBJCXX)
-     target_link_libraries(${lib_name})
-   endforeach()
-
-   list(APPEND MM_SRCS
-     buffer_convert
-     common
-   )
-   foreach(lib_name ${MM_SRCS})
-     add_library("${lib_name}" STATIC ${METAL_DELEGATE_PATH}${lib_name}.mm)
-     target_include_directories("${lib_name}" PUBLIC
-       ${CMAKE_BINARY_DIR}/abseil-cpp
-       ${CMAKE_BINARY_DIR}/flatbuffers/include
-     )
-     target_link_libraries(${lib_name})
-   endforeach()
 endif()
   list(APPEND TFLITE_TARGET_PUBLIC_OPTIONS "-DCL_DELEGATE_NO_GL" "-DEGL_NO_X11")
   list(APPEND TFLITE_TARGET_DEPENDENCIES
diff --git a/tensorflow/lite/fuzzing/BUILD b/tensorflow/lite/fuzzing/BUILD
new file mode 100644
index 00000000000000..223e7a35c0898b
--- /dev/null
+++ b/tensorflow/lite/fuzzing/BUILD
@@ -0,0 +1,22 @@
+# Fuzzing harnesses for the TensorFlow Lite runtime.
+
+load(
+    "//tensorflow/security/fuzzing:tf_fuzzing.bzl",
+    "tf_cc_fuzz_test",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = ["//visibility:private"],
+    licenses = ["notice"],
+)
+
+tf_cc_fuzz_test(
+    name = "interpreter_fuzz",
+    srcs = ["interpreter_fuzz.cc"],
+    deps = [
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/kernels:builtin_ops",
+    ],
+)
diff --git a/tensorflow/lite/fuzzing/interpreter_fuzz.cc b/tensorflow/lite/fuzzing/interpreter_fuzz.cc
new file mode 100644
index 00000000000000..785090ace761bd
--- /dev/null
+++ b/tensorflow/lite/fuzzing/interpreter_fuzz.cc
@@ -0,0 +1,99 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+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.
+==============================================================================*/
+
+// Fuzzes the TensorFlow Lite interpreter end to end: an arbitrary buffer is
+// verified as a flatbuffer model, an interpreter is built for it, tensors are
+// allocated, inputs are filled deterministically, and the graph is invoked.
+//
+// This exercises the builtin kernel implementations, the arena planner and the
+// shape-propagation paths, none of which previously had OSS-Fuzz coverage.
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <memory>
+#include <string>
+
+#include "fuzztest/fuzztest.h"
+#include "tensorflow/lite/core/interpreter.h"
+#include "tensorflow/lite/core/interpreter_builder.h"
+#include "tensorflow/lite/core/model_builder.h"
+#include "tensorflow/lite/kernels/register.h"
+
+namespace tflite {
+namespace fuzzing {
+namespace {
+
+// Keep the fuzzer inside the OSS-Fuzz memory budget. Models and arenas larger
+// than this are not interesting: they exercise the allocator, not the kernels.
+constexpr size_t kMaxModelBytes = 1 << 20;  // 1 MiB
+constexpr size_t kMaxArenaBytes = 1 << 26;  // 64 MiB
+
+void FuzzInterpreter(const std::string& model_bytes) {
+  if (model_bytes.size() < 8 || model_bytes.size() > kMaxModelBytes) {
+    return;
+  }
+
+  // VerifyAndBuildFromBuffer applies tflite::VerifyModelBuffer first, so
+  // structurally invalid buffers are rejected cheaply.
+  std::unique_ptr<FlatBufferModel> model =
+      FlatBufferModel::VerifyAndBuildFromBuffer(model_bytes.data(),
+                                                model_bytes.size());
+  if (model == nullptr) {
+    return;
+  }
+
+  // Delegates are excluded so the fuzzer exercises the reference and optimized
+  // CPU kernels rather than a delegate's own implementation.
+  ops::builtin::BuiltinOpResolverWithoutDefaultDelegates resolver;
+  std::unique_ptr<Interpreter> interpreter;
+  if (InterpreterBuilder(*model, resolver)(&interpreter) != kTfLiteOk ||
+      interpreter == nullptr) {
+    return;
+  }
+
+  if (interpreter->AllocateTensors() != kTfLiteOk) {
+    return;
+  }
+
+  size_t total_bytes = 0;
+  for (const int tensor_index : interpreter->inputs()) {
+    TfLiteTensor* tensor = interpreter->tensor(tensor_index);
+    if (tensor == nullptr || tensor->data.raw == nullptr) {
+      continue;
+    }
+    // String tensors own a dynamic buffer with its own layout; writing raw
+    // bytes into it would corrupt the interpreter rather than the kernel under
+    // test.
+    if (tensor->type == kTfLiteString || tensor->type == kTfLiteResource ||
+        tensor->type == kTfLiteVariant) {
+      return;
+    }
+    // Check before accumulating so the sum itself cannot wrap.
+    if (tensor->bytes > kMaxArenaBytes ||
+        total_bytes > kMaxArenaBytes - tensor->bytes) {
+      return;
+    }
+    total_bytes += tensor->bytes;
+    std::memset(tensor->data.raw, 1, tensor->bytes);
+  }
+
+  interpreter->Invoke();
+}
+FUZZ_TEST(TfLiteFuzz, FuzzInterpreter);
+
+}  // namespace
+}  // namespace fuzzing
+}  // namespace tflite
diff --git a/tensorflow/lite/kernels/add_test.cc b/tensorflow/lite/kernels/add_test.cc
index 0469a7d0ec7c92..92249345046f7b 100644
--- a/tensorflow/lite/kernels/add_test.cc
+++ b/tensorflow/lite/kernels/add_test.cc
@@ -17,8 +17,10 @@ limitations under the License.
 
 #include <algorithm>
 #include <array>
+#include <limits>
 #include <numeric>
 #include <random>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
@@ -736,6 +738,37 @@ TYPED_TEST(IntegerAddOpTest, Int32MultiDimBroadcast) {
   EXPECT_THAT(m.GetOutput<TypeParam>(), ElementsAreArray({4, 6, 7, 9}));
 }
 
+TYPED_TEST(IntegerAddOpTest, OverflowWrapping) {
+  if (std::is_same<TypeParam, int32_t>::value ||
+      std::is_same<TypeParam, int64_t>::value) {
+    IntegerAddOpModel m(GetTensorType<TypeParam>(), {1, 2}, {1, 2},
+                        ActivationFunctionType_NONE);
+    m.PopulateTensor<TypeParam>(m.input1(),
+                                {std::numeric_limits<TypeParam>::max(), 5});
+    m.PopulateTensor<TypeParam>(m.input2(), {1, 5});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutput<TypeParam>(),
+                ElementsAreArray<TypeParam>(
+                    {std::numeric_limits<TypeParam>::min(), 10}));
+  }
+}
+
+TYPED_TEST(IntegerAddOpTest, OverflowWrappingBroadcast) {
+  if (std::is_same<TypeParam, int32_t>::value ||
+      std::is_same<TypeParam, int64_t>::value) {
+    IntegerAddOpModel m(GetTensorType<TypeParam>(), {1, 2}, {2, 1},
+                        ActivationFunctionType_NONE);
+    m.PopulateTensor<TypeParam>(m.input1(),
+                                {std::numeric_limits<TypeParam>::max(), 5});
+    m.PopulateTensor<TypeParam>(m.input2(), {1, 1});
+    ASSERT_EQ(m.Invoke(), kTfLiteOk);
+    EXPECT_THAT(m.GetOutput<TypeParam>(),
+                ElementsAreArray<TypeParam>(
+                    {std::numeric_limits<TypeParam>::min(), 6,
+                     std::numeric_limits<TypeParam>::min(), 6}));
+  }
+}
+
 template <typename QuantizedType>
 void TestQuantizedBroadcast(QuantizedAddOpModel& m,
                             const std::vector<int>& input1_shape,
diff --git a/tensorflow/lite/kernels/internal/optimized/optimized_ops.h b/tensorflow/lite/kernels/internal/optimized/optimized_ops.h
index 1940086a9ae9e0..dab0bd8efda225 100644
--- a/tensorflow/lite/kernels/internal/optimized/optimized_ops.h
+++ b/tensorflow/lite/kernels/internal/optimized/optimized_ops.h
@@ -1828,20 +1828,32 @@ inline typename std::enable_if<is_int32_or_int64<T>::value, void>::type Add(
   auto input1_map = MapAsVector(input1_data, input1_shape);
   auto input2_map = MapAsVector(input2_data, input2_shape);
   auto output_map = MapAsVector(output_data, output_shape);
+  // The element-wise sum is performed in the unsigned domain so that any
+  // overflow wraps in a well-defined manner instead of triggering
+  // signed-integer-overflow UB inside Eigen's expression evaluators. The
+  // wrapped result is bit-identical to the previous two's-complement behavior
+  // and is then clamped back in the signed domain by the activation min/max.
+  using UnsignedT = typename std::make_unsigned<T>::type;
   if (input1_shape == input2_shape) {
-    output_map.array() = (input1_map.array() + input2_map.array())
+    output_map.array() = (input1_map.array().template cast<UnsignedT>() +
+                          input2_map.array().template cast<UnsignedT>())
+                             .template cast<T>()
                              .cwiseMax(activation_min)
                              .cwiseMin(activation_max);
   } else if (input2_shape.FlatSize() == 1) {
-    auto scalar = input2_data[0];
-    output_map.array() = (input1_map.array() + scalar)
-                             .cwiseMax(activation_min)
-                             .cwiseMin(activation_max);
+    UnsignedT scalar = static_cast<UnsignedT>(input2_data[0]);
+    output_map.array() =
+        (input1_map.array().template cast<UnsignedT>() + scalar)
+            .template cast<T>()
+            .cwiseMax(activation_min)
+            .cwiseMin(activation_max);
   } else if (input1_shape.FlatSize() == 1) {
-    auto scalar = input1_data[0];
-    output_map.array() = (scalar + input2_map.array())
-                             .cwiseMax(activation_min)
-                             .cwiseMin(activation_max);
+    UnsignedT scalar = static_cast<UnsignedT>(input1_data[0]);
+    output_map.array() =
+        (scalar + input2_map.array().template cast<UnsignedT>())
+            .template cast<T>()
+            .cwiseMax(activation_min)
+            .cwiseMin(activation_max);
   } else {
     reference_ops::BroadcastAdd6DSlow<T>(params, input1_shape, input1_data,
                                          input2_shape, input2_data,
diff --git a/tensorflow/lite/kernels/internal/reference/add.h b/tensorflow/lite/kernels/internal/reference/add.h
index 198d576c4863c0..fe031644da1403 100644
--- a/tensorflow/lite/kernels/internal/reference/add.h
+++ b/tensorflow/lite/kernels/internal/reference/add.h
@@ -41,7 +41,8 @@ inline void Add(const ArithmeticParams& params,
       MatchingElementsSize(input1_shape, input2_shape, output_shape);
   for (int i = 0; i < flat_size; ++i) {
     output_data[i] = ActivationFunctionWithMinMax<T>(
-        input1_data[i] + input2_data[i], activation_min, activation_max);
+        WrappingAdd(input1_data[i], input2_data[i]), activation_min,
+        activation_max);
   }
 }
 
@@ -280,7 +281,7 @@ BroadcastAdd6DSlow(const ArithmeticParams& params,
   T activation_min, activation_max;
   GetActivationParams(params, &activation_min, &activation_max);
   auto op = [activation_min, activation_max](T a, T b) {
-    return ActivationFunctionWithMinMax<T>(a + b, activation_min,
+    return ActivationFunctionWithMinMax<T>(WrappingAdd(a, b), activation_min,
                                            activation_max);
   };
   BroadcastBinaryOpSimple(input1_shape, input1_data, input2_shape, input2_data,
diff --git a/tensorflow/lite/kernels/slice.cc b/tensorflow/lite/kernels/slice.cc
index 86dca38bb7c427..597c14f585d25c 100644
--- a/tensorflow/lite/kernels/slice.cc
+++ b/tensorflow/lite/kernels/slice.cc
@@ -145,7 +145,16 @@ TfLiteStatus Prepare(TfLiteContext* context, TfLiteNode* node) {
   TF_LITE_ENSURE_EQ(context, NumElements(begin), NumElements(size));
   // If the shape of output is fully specified then resize even if
   // the input shape is not staticly defined.
-  if (!HasUnspecifiedDimension(output) && ShapeHasRank(output->dims)) {
+  //
+  // A fully specified output shape does not imply the slice is in bounds when
+  // the input extent is only known at run time. Taking this path leaves the
+  // output static, so Eval() -- which only re-runs ResizeOutputShape() for a
+  // dynamic output -- never reaches CalculateOutputShapeVector(), the one place
+  // `begin` and `size` are checked against the input. Fall through when the
+  // input has an unspecified dimension so the output is marked dynamic and the
+  // bounds are validated against the actual extent on every invocation.
+  if (ShapeHasRank(output->dims) && !HasUnspecifiedDimension(output) &&
+      ShapeHasRank(input->dims) && !HasUnspecifiedDimension(input)) {
     return kTfLiteOk;
   }
   // Postpone allocation of output if any of the indexing tensors is not
diff --git a/tensorflow/lite/kernels/slice_test.cc b/tensorflow/lite/kernels/slice_test.cc
index 09b68642fff848..e3b55f6116e32c 100644
--- a/tensorflow/lite/kernels/slice_test.cc
+++ b/tensorflow/lite/kernels/slice_test.cc
@@ -102,6 +102,73 @@ class SliceOpModel : public SingleOpModel {
 
 class SliceOpTest : public ::testing::TestWithParam<TestType> {};
 
+// Model with a dynamic input dimension and a statically shaped output. The
+// suite name is deliberately distinct from SliceOpTest, which is a TEST_P
+// fixture -- gtest rejects a suite that mixes TEST and TEST_P.
+class DynamicInputSliceOpModel : public SingleOpModel {
+ public:
+  DynamicInputSliceOpModel(TensorData input_data,
+                           std::initializer_list<int> begin_shape,
+                           std::initializer_list<int32_t> begin_data,
+                           std::initializer_list<int> size_shape,
+                           std::initializer_list<int32_t> size_data,
+                           TensorData output_data) {
+    input_ = AddInput(input_data);
+    begin_ = AddConstInput(TensorType_INT32, begin_data, begin_shape);
+    size_ = AddConstInput(TensorType_INT32, size_data, size_shape);
+    output_ = AddOutput(output_data);
+    SetBuiltinOp(BuiltinOperator_SLICE, BuiltinOptions_SliceOptions,
+                 CreateSliceOptions(builder_).Union());
+    // Delegates are bypassed: a delegate that claims the SLICE node would set
+    // the output allocation type itself, so the assertions below would no
+    // longer describe the built-in CPU kernel.
+    BuildInterpreter({input_data.shape, begin_shape, size_shape},
+                     /*num_threads=*/-1, /*allow_fp32_relax_to_fp16=*/false,
+                     /*apply_delegate=*/false);
+  }
+
+  void SetInput(std::initializer_list<float> data) {
+    PopulateTensor<float>(input_, data);
+  }
+  std::vector<float> GetOutput() { return ExtractVector<float>(output_); }
+  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }
+  const TfLiteTensor* GetOutputTensor() {
+    return interpreter_->tensor(output_);
+  }
+
+ private:
+  int input_;
+  int begin_;
+  int size_;
+  int output_;
+};
+
+// A dynamic input dimension must force the output dynamic so the bounds are
+// re-validated in Eval(), even though the declared output shape is static.
+TEST(SliceOpDynamicInputTest, DynamicInputStaticOutputValid) {
+  TensorData input_data(TensorType_FLOAT32, {1, 8});
+  input_data.shape_signature = {1, -1};
+  TensorData output_data(TensorType_FLOAT32, {1, 4});
+
+  DynamicInputSliceOpModel m(input_data, {2}, {0, 2}, {2}, {1, 4}, output_data);
+  m.SetInput({1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0});
+  ASSERT_EQ(m.Invoke(), kTfLiteOk);
+  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1, 4}));
+  EXPECT_THAT(m.GetOutput(), ElementsAreArray({3.0, 4.0, 5.0, 6.0}));
+  EXPECT_EQ(m.GetOutputTensor()->allocation_type, kTfLiteDynamic);
+}
+
+// The same path must reject a window that does not fit the actual extent.
+TEST(SliceOpDynamicInputTest, DynamicInputStaticOutputOutOfBounds) {
+  TensorData input_data(TensorType_FLOAT32, {1, 4});
+  input_data.shape_signature = {1, -1};
+  TensorData output_data(TensorType_FLOAT32, {1, 8});
+
+  DynamicInputSliceOpModel m(input_data, {2}, {0, 0}, {2}, {1, 8}, output_data);
+  m.SetInput({1.0, 2.0, 3.0, 4.0});
+  EXPECT_EQ(m.Invoke(), kTfLiteError);
+}
+
 TEST_P(SliceOpTest, In1D) {
   SliceOpModel<float, int32_t> m({4}, {1}, {1}, {1}, {2}, TensorType_INT32,
                                  TensorType_FLOAT32, GetParam());
diff --git a/tensorflow/lite/kernels/transpose.cc b/tensorflow/lite/kernels/transpose.cc
index 0b1f2b783b05bd..94bdac9439a21b 100644
--- a/tensorflow/lite/kernels/transpose.cc
+++ b/tensorflow/lite/kernels/transpose.cc
@@ -50,12 +50,26 @@ TfLiteStatus ResizeOutputTensor(TfLiteContext* context,
   // Ensure validity of the permutations tensor as a 1D tensor.
   TF_LITE_ENSURE_EQ(context, NumDimensions(op_context->perm), 1);
   TF_LITE_ENSURE_EQ(context, op_context->perm->dims->data[0], dims);
+  // `perm` must be a permutation of [0, dims), not merely a set of in-range
+  // values: the output shape and the element offsets are both derived from it,
+  // and a repeated entry makes them disagree with the input extent.
+  // `dims` is bounded by kTransposeMaxDimensions, which Prepare() enforces
+  // before this function is reachable, so a 64-bit mask is sufficient and
+  // avoids a heap allocation in the kernel.
+  static_assert(kTransposeMaxDimensions <= 64,
+                "Permutation bitmask assumes at most 64 dimensions.");
+  uint64_t seen = 0;
   for (int idx = 0; idx < dims; ++idx) {
     TF_LITE_ENSURE_MSG(context,
                        (perm_data[idx] >= -dims && perm_data[idx] < dims),
                        "Transpose op permutations array is out of bounds.");
     new_perm_data[idx] = perm_data[idx];
     if (new_perm_data[idx] < 0) new_perm_data[idx] += dims;
+    const uint64_t bit = uint64_t{1} << new_perm_data[idx];
+    TF_LITE_ENSURE_MSG(
+        context, (seen & bit) == 0,
+        "Transpose op permutations array must not contain duplicate values.");
+    seen |= bit;
   }
 
   // Determine size of output tensor.
diff --git a/tensorflow/lite/kernels/transpose_test.cc b/tensorflow/lite/kernels/transpose_test.cc
index 601e8bf4355540..11930f342ce91a 100644
--- a/tensorflow/lite/kernels/transpose_test.cc
+++ b/tensorflow/lite/kernels/transpose_test.cc
@@ -193,6 +193,23 @@ TEST(TransposeTest, TestPermOutOfBounds) {
   EXPECT_DEATH(TransposeOpConstModel({1, 3, 3, 1}, {4}, {0, 1, 2, 4}),
                "Transpose op permutations array is out of bounds.");
 }
+
+// `perm` must be a permutation of [0, dims). A repeated entry passes the range
+// check but makes the derived output shape disagree with the input extent, so
+// the kernel reads outside the input tensor.
+TEST(TransposeTest, TestPermDuplicateValues) {
+  EXPECT_DEATH(
+      TransposeOpConstModel({1, 3, 3, 1}, {4}, {0, 1, 2, 2}),
+      "Transpose op permutations array must not contain duplicate values.");
+}
+
+// Duplicates must also be rejected after negative entries are normalised:
+// on a rank-2 input {0, -2} normalises to {0, 0}.
+TEST(TransposeTest, TestPermDuplicateValuesAfterNegativeNormalization) {
+  EXPECT_DEATH(
+      TransposeOpConstModel({2, 3}, {2}, {0, -2}),
+      "Transpose op permutations array must not contain duplicate values.");
+}
 #endif
 
 TEST(TransposeTest, TestInt41DInputConstTensor) {
diff --git a/tensorflow/lite/python/lite.py b/tensorflow/lite/python/lite.py
index 470c58779a768b..bd9af24daa0491 100644
--- a/tensorflow/lite/python/lite.py
+++ b/tensorflow/lite/python/lite.py
@@ -1584,18 +1584,27 @@ def convert(self):
           graph_def, input_tensors, output_tensors
       )
 
-    trackable_obj = _load(self.saved_model_dir, self._saved_model_tags)
-    if trackable_obj is None:
-      self._debug_info = _get_debug_info(
-          _build_debug_info_func(self._funcs[0].graph), graph_def
-      )
-    else:
-      self._debug_info = _get_debug_info(
-          _convert_debug_info_func(trackable_obj.graph_debug_info),
-          graph_def,
-      )
-
-    del trackable_obj
+    # Read debug info directly from the SavedModel directory instead of loading
+    # the full model with _load(). Loading the model allocates all variable
+    # tensors (~25MB for DenseNet121) plus registers function defs in the
+    # EagerContext, and these can leak across convert() calls due to reference
+    # cycles that are slow to collect.  The debug info proto is already written
+    # to disk alongside saved_model.pb and can be read cheaply.
+    #
+    # Note on adjust_debug_info_func_names: _convert_debug_info_func() (in
+    # util.py) ignores its original_nodes argument entirely and returns the
+    # saved_debug_info proto unchanged, so the function-name adjustment that
+    # _load() applies via adjust_debug_info_func_names() has no effect on this
+    # code path.  The on-disk proto already contains the names as written during
+    # save(), which is what the TFLite C++ pipeline uses for source-location
+    # tracking.
+    _, saved_debug_info = _parse_saved_model_with_debug_info(
+        self.saved_model_dir
+    )
+    self._debug_info = _get_debug_info(
+        _convert_debug_info_func(saved_debug_info),
+        graph_def,
+    )
     return self._convert_from_saved_model(graph_def)
 
 
diff --git a/tensorflow/lite/python/lite_v2_test.py b/tensorflow/lite/python/lite_v2_test.py
index 025a3a3b914105..ae73b2046f6731 100644
--- a/tensorflow/lite/python/lite_v2_test.py
+++ b/tensorflow/lite/python/lite_v2_test.py
@@ -3169,6 +3169,47 @@ def testQDQConversionMode(self, mode):
       model = converter.convert()
       self.assertIsNotNone(model)
 
+  @test_util.run_v2_only
+  def testConvertDoesNotCallLoad(self):
+    """Regression test for https://github.com/tensorflow/tensorflow/issues/122598.
+
+    TFLiteSavedModelConverterV2.convert() used to call _load() internally just
+    to obtain graph_debug_info. For large models like DenseNet121, _load()
+    allocates ~25 MB of variable tensors and registers function defs in TF's
+    EagerContext. Reference cycles in the loaded object caused this memory to
+    leak across convert() calls (~22 MB/iter). The fix reads debug info directly
+    from the SavedModel directory via _parse_saved_model_with_debug_info()
+    instead of loading the full model.
+    """
+    from unittest import mock
+
+    root = autotrackable.AutoTrackable()
+    root.f = tf.function(lambda x: x * 2.0)
+    to_save = root.f.get_concrete_function(tf.TensorSpec([10], tf.float32))
+    save_dir = os.path.join(self.get_temp_dir(), 'saved_model_no_load')
+    save.save(root, save_dir, to_save)
+
+    converter = lite.TFLiteConverterV2.from_saved_model(save_dir)
+    # Patch _load at the lite module level to detect if it's called during
+    # convert(). The converter should NOT call _load() — it should read debug
+    # info directly from disk instead.
+    with mock.patch.object(lite, '_load', wraps=lite._load) as mock_load:
+      tflite_model = converter.convert()
+      mock_load.assert_not_called()
+
+    self.assertIsNotNone(tflite_model)
+    # Verify the converted model is runnable.
+    interp = interpreter.Interpreter(model_content=tflite_model)
+    interp.allocate_tensors()
+    input_details = interp.get_input_details()
+    output_details = interp.get_output_details()
+    interp.set_tensor(
+        input_details[0]['index'], np.ones([10], dtype=np.float32)
+    )
+    interp.invoke()
+    output = interp.get_tensor(output_details[0]['index'])
+    np.testing.assert_array_almost_equal(output, np.full([10], 2.0))
+
 
 class FromKerasModelTest(lite_v2_test_util.ModelTest):
 
diff --git a/third_party/xla/.github/workflows/rocm_ci.yml b/third_party/xla/.github/workflows/rocm_ci.yml
index 1e08a8013c9b0e..a722371d30ade1 100644
--- a/third_party/xla/.github/workflows/rocm_ci.yml
+++ b/third_party/xla/.github/workflows/rocm_ci.yml
@@ -42,12 +42,17 @@ jobs:
     runs-on: ubuntu-latest
     outputs:
       docker-image: ${{ steps.out.outputs.docker-image }}
+      rocm-distro-url: ${{ steps.out.outputs.rocm-distro-url }}
+      rocm-distro-hash: ${{ steps.out.outputs.rocm-distro-hash }}
     steps:
       - id: out
         shell: bash
         run: |
-          # hermetic llvm based on  ghcr.io/rocm/jax-ubu22.rocm7.2.4:latest
-          echo "docker-image=ghcr.io/rocm/jax-ubu22.rocm7.2.4@sha256:248db525342fb7438ccbc6f5f1396ec82c783f7e0ae5667cb698edcddc8fe2fd" >> "$GITHUB_OUTPUT"
+          {
+            echo "docker-image=ghcr.io/rocm/jax-build-ubu24.rocmless@sha256:68854ed30b800d3e8c390f6e651721e2bead58830da1c59e1b828354eb8f077f"
+            echo "rocm-distro-url=https://stable.repo.amd.com/rocm/core/tarball/therock-dist-linux-multiarch-10.0.0.tar.gz"
+            echo "rocm-distro-hash=1c5e807875d26a2470ecc7323daa5b5b9009208a55c3290ac255a909cde15fc6"
+          } >> "$GITHUB_OUTPUT"
 
   jax:
     needs: rocm-config
@@ -59,6 +64,8 @@ jobs:
     continue-on-error: true
     env:
       DOCKER_IMAGE: ${{ needs.rocm-config.outputs.docker-image }}
+      ROCM_DISTRO_URL: ${{ needs.rocm-config.outputs.rocm-distro-url }}
+      ROCM_DISTRO_HASH: ${{ needs.rocm-config.outputs.rocm-distro-hash }}
       # Unique per run so a crashed run cannot collide with this one on a pooled node.
       CONTAINER_NAME: jax-${{ github.run_id }}-${{ github.run_attempt }}
       # Stable label so leftover containers from crashed runs can be swept on a pooled node.
@@ -150,6 +157,10 @@ jobs:
             --repo_env=HERMETIC_PYTHON_VERSION=3.14 \
             --repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \
             --jobs=20 \
+            --repo_env=ROCM_PATH="" \
+            --repo_env=ROCM_DISTRO_URL="${ROCM_DISTRO_URL}" \
+            --repo_env=ROCM_DISTRO_HASH="${ROCM_DISTRO_HASH}" \
+            --repo_env=SYSROOT_DIST=linux_glibc_2_31 \
             --bes_keywords=jax \
             --bes_keywords=upstream \
             --bes_keywords=gpu \
@@ -166,7 +177,10 @@ jobs:
     timeout-minutes: 220
     env:
       DOCKER_IMAGE: ${{ needs.rocm-config.outputs.docker-image }}
+      ROCM_DISTRO_URL: ${{ needs.rocm-config.outputs.rocm-distro-url }}
+      ROCM_DISTRO_HASH: ${{ needs.rocm-config.outputs.rocm-distro-hash }}
       EXECUTE_CI_BUILD_URL: https://raw.githubusercontent.com/ROCm/xla/refs/heads/${{ inputs.rocm_xla_branch || 'rocm-dev-infra' }}/build_tools/rocm/execute_ci_build_upstream.sh
+      GLOBAL_SYMBOL_VERSION_LDS_URL: https://raw.githubusercontent.com/ROCm/xla/refs/heads/${{ inputs.rocm_xla_branch || 'rocm-dev-infra' }}/build_tools/rocm/global_symbol_version.lds
       # Unique per run so a crashed run cannot collide with this one on a pooled node.
       CONTAINER_NAME: xla-${{ github.run_id }}-${{ github.run_attempt }}
       # Stable label so leftover containers from crashed runs can be swept on a pooled node.
@@ -181,10 +195,21 @@ jobs:
           repository: openxla/xla
           persist-credentials: false
 
-      - name: Fetch ROCm XLA Upstream Script from ROCm/xla
+      - name: Fetch ROCm XLA Upstream Scripts from ROCm/xla
         run: |
           wget -O build_tools/rocm/execute_ci_build_upstream.sh "${EXECUTE_CI_BUILD_URL}"
           chmod +x build_tools/rocm/execute_ci_build_upstream.sh
+          wget -O build_tools/rocm/global_symbol_version.lds "${GLOBAL_SYMBOL_VERSION_LDS_URL}"
+          LDS_SHA=$(sha256sum build_tools/rocm/global_symbol_version.lds | cut -d' ' -f1)
+          # Workaround for the symbol clash with ROCm's libLLVM.so
+          {
+            echo "LLVM_SYMBOL_CLASH_WAR<<EOF"
+            echo "--dynamic_mode=default"
+            echo "--linkopt=-Wl,--version-script,/work/build_tools/rocm/global_symbol_version.lds"
+            echo "--sandbox_add_mount_pair=/work/build_tools/rocm/global_symbol_version.lds"
+            echo "--action_env=global_symbol_version_lds_CACHEBUSTER=${LDS_SHA}"
+            echo "EOF"
+          } >> "$GITHUB_ENV"
 
       - *start_container
 
@@ -194,16 +219,20 @@ jobs:
       - name: Test XLA [single_gpu]
         timeout-minutes: 120
         run: |
+          # LLVM_SYMBOL_CLASH_WAR is intentionally unquoted so it splits into one flag per line.
+          # shellcheck disable=SC2086
           docker exec "${CONTAINER_NAME}" build_tools/rocm/execute_ci_build_upstream.sh \
-            --config=rocm_ci \
+            --config=rocm_ci_hermetic \
             --config=rocm_rbe \
             --config=ci_single_gpu \
-            --local_test_jobs=1 \
             --repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \
-            --repo_env=ROCM_PATH="/opt/rocm" \
+            --repo_env=ROCM_PATH="" \
+            --repo_env=ROCM_DISTRO_URL="${ROCM_DISTRO_URL}" \
+            --repo_env=ROCM_DISTRO_HASH="${ROCM_DISTRO_HASH}" \
             --local_test_jobs=1 \
             --internal_spawn_scheduler \
             --strategy=TestRunner=dynamic \
+            ${LLVM_SYMBOL_CLASH_WAR} \
             --bes_keywords=xla \
             --bes_keywords=upstream \
             --bes_keywords=gpu \
@@ -212,13 +241,18 @@ jobs:
       - name: Test XLA [rocm_cpu]
         timeout-minutes: 80
         run: |
+          # LLVM_SYMBOL_CLASH_WAR is intentionally unquoted so it splits into one flag per line.
+          # shellcheck disable=SC2086
           docker exec "${CONTAINER_NAME}" build_tools/rocm/execute_ci_build_upstream.sh \
-            --config=rocm_ci \
+            --config=rocm_ci_hermetic \
             --config=rocm_rbe \
             --config=ci_rocm_cpu \
             --local_test_jobs=200 \
             --repo_env=TF_ROCM_RBE_DOCKER_IMAGE="${DOCKER_IMAGE}" \
-            --repo_env=ROCM_PATH="/opt/rocm" \
+            --repo_env=ROCM_PATH="" \
+            --repo_env=ROCM_DISTRO_URL="${ROCM_DISTRO_URL}" \
+            --repo_env=ROCM_DISTRO_HASH="${ROCM_DISTRO_HASH}" \
+            ${LLVM_SYMBOL_CLASH_WAR} \
             --bes_keywords=xla \
             --bes_keywords=upstream \
             --bes_keywords=cpu \
diff --git a/third_party/xla/tensorflow.bazelrc b/third_party/xla/tensorflow.bazelrc
index d6ceb7b9829bc5..054ade6bdf69c5 100644
--- a/third_party/xla/tensorflow.bazelrc
+++ b/third_party/xla/tensorflow.bazelrc
@@ -292,20 +292,20 @@ common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_rocm=True
 common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_cuda=False
 common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_sycl=False
 common:rocm_clang_hermetic --@rules_ml_toolchain//common:enable_hermetic_cc=True
-common:rocm_clang_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic
+common:rocm_clang_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic --config=workspace
 common:rocm_clang_hermetic --strategy=CppLink=local
 common:rocm_clang_hermetic --@rules_ml_toolchain//common:static_libcxx=False
 
 common:rocm --config=rocm_clang_hermetic
 common:rocm_ci --config=rocm
-common:rocm_ci --@local_config_rocm//rocm:rocm_path_type=hermetic
+common:rocm_ci --@local_config_rocm//rocm:rocm_path_type=hermetic --config=workspace
 
 common:rocm_ci_hermetic --dynamic_mode=off
 common:rocm_ci_hermetic --config=rocm_clang_hermetic
 common:rocm_ci_hermetic --repo_env=TF_ROCM_AMDGPU_TARGETS="gfx908,gfx90a"
 common:rocm_ci_hermetic --repo_env=ROCM_DISTRO_VERSION="rocm_7.13.0_gfx90a"
 common:rocm_ci_hermetic --repo_env=SYSROOT_DIST=linux_glibc_2_31
-common:rocm_ci_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic
+common:rocm_ci_hermetic --@local_config_rocm//rocm:rocm_path_type=hermetic --config=workspace
 
 # This config option is used for SYCL as GPU backend.
 # SYCL Configuration (non-hermetic)
diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/non_identity_layout_unit_dim.hlo b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/non_identity_layout_unit_dim.hlo
new file mode 100644
index 00000000000000..e93415f2ad4d6b
--- /dev/null
+++ b/third_party/xla/xla/backends/cpu/codegen/tiled/tests/transpose/non_identity_layout_unit_dim.hlo
@@ -0,0 +1,27 @@
+// RUN: fusion_to_xtile %s | fusion_compiler_opt \
+// RUN:   --xtile-cpu-new-xtile-to-vector --xtile-cpu-new-vector-to-llvm \
+// RUN:   | FileCheck %s --check-prefix=NEW-VECTOR-LLVM
+// RUN: test_correctness %s --xla_cpu_use_new_xtile_lowering=true
+
+fusion {
+  param_0 = f32[2,1,2,1]{3,1,0,2} parameter(0)
+  copy_0 = f32[2,1,2,1]{3,2,1,0} copy(param_0)
+  transpose = f32[2,2,1,1]{3,0,2,1} transpose(copy_0), dimensions={2,0,1,3}
+  copy_1 = f32[2,2,1,1]{3,2,1,0} copy(transpose)
+  ROOT bitcast = f32[2,2]{1,0} bitcast(copy_1)
+}
+
+ENTRY main {
+  param_0 = f32[2,1,2,1]{3,1,0,2} parameter(0)
+  ROOT fusion = f32[2,2]{1,0} fusion(param_0), kind=kLoop, calls=fusion,
+    backend_config={
+      "fusion_backend_config":{
+        "block_level_fusion_config":{
+          "output_tiles":[{"sizes":["2", "2"]}],
+          "num_ctas":"1"
+        }
+      }
+    }
+}
+// NEW-VECTOR-LLVM-LABEL: llvm.func internal @wrapped_fusion_impl
+// NEW-VECTOR-LLVM-NOT: llvm.alloc
diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/drop_vector_unit_dims_pass.cc b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/drop_vector_unit_dims_pass.cc
index 798e110957cc43..7663a570e87dc5 100644
--- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/drop_vector_unit_dims_pass.cc
+++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/drop_vector_unit_dims_pass.cc
@@ -38,8 +38,6 @@ class DropVectorUnitDimsPass
     mlir::RewritePatternSet patterns(context);
     mlir::vector::populateCastAwayVectorLeadingOneDimPatterns(patterns);
     mlir::vector::populateDropUnitDimWithShapeCastPatterns(patterns);
-    mlir::vector::populateDropInnerMostUnitDimsXferOpPatterns(patterns);
-    mlir::vector::populateVectorTransferDropUnitDimsPatterns(patterns);
     if (mlir::failed(
             mlir::applyPatternsGreedily(getOperation(), std::move(patterns)))) {
       signalPassFailure();
diff --git a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/drop_vector_unit_dims_pass.mlir b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/drop_vector_unit_dims_pass.mlir
index 569e8ec080d24b..35e4a68845bc14 100644
--- a/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/drop_vector_unit_dims_pass.mlir
+++ b/third_party/xla/xla/backends/cpu/codegen/tiled/transforms/tests/drop_vector_unit_dims_pass.mlir
@@ -65,3 +65,17 @@ func.func @no_unit_dims(%arg0: vector<16xf32>, %arg1: vector<16xf32>) -> vector<
 // CHECK-SAME:     %[[ARG0:.*]]: vector<16xf32>, %[[ARG1:.*]]: vector<16xf32>) -> vector<16xf32> {
 // CHECK-NOT: vector.shape_cast
 // CHECK:     arith.addf %[[ARG0]], %[[ARG1]] : vector<16xf32>
+
+// -----
+
+func.func @transfer_read_non_identity_layout_unit_dim(
+    %arg0: memref<2x1x2x1xf32, #xtile.layout<[3, 1, 0, 2]>>, %c0: index) -> vector<2x1x2x1xf32> {
+  %pad = arith.constant 0.0 : f32
+  %0 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %pad {in_bounds = [true, true, true, true]} : memref<2x1x2x1xf32, #xtile.layout<[3, 1, 0, 2]>>, vector<2x1x2x1xf32>
+  return %0 : vector<2x1x2x1xf32>
+}
+// CHECK-LABEL: func.func @transfer_read_non_identity_layout_unit_dim(
+// CHECK-SAME:     %[[ARG0:.*]]: memref<2x1x2x1xf32, #xtile.layout<[3, 1, 0, 2]>>, %[[C0:.*]]: index) -> vector<2x1x2x1xf32> {
+// CHECK-NOT:  memref.subview
+// CHECK:      vector.transfer_read %[[ARG0]]
+
diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc
index c9025869505fc8..d20b178c722844 100644
--- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler.cc
@@ -512,26 +512,61 @@ class ConvDimensionAdapter {
                                 dnums_for_layout};
   }
 
+  int64_t HloDimToCudnnDim(int64_t hlo_dim) const {
+    if (hlo_dim == dums_.output_batch_dimension()) {
+      return 0;  // Batch (N)
+    }
+    if (hlo_dim == dums_.output_feature_dimension()) {
+      return 1;  // Feature / Channel (C)
+    }
+    int64_t dummy_spatial_dims =
+        std::max<int64_t>(0, 2 - dums_.output_spatial_dimensions_size());
+    for (int i = 0; i < dums_.output_spatial_dimensions_size(); ++i) {
+      if (hlo_dim == dums_.output_spatial_dimensions(i)) {
+        return 2 + dummy_spatial_dims + i;  // Spatial dimensions (H, W, ...)
+      }
+    }
+    return -1;
+  }
+
   std::optional<Result> DimensionsAndStrides(const HloInstruction& hlo) {
+    int64_t spatial_dims =
+        std::max<int64_t>(2, dums_.input_spatial_dimensions_size());
+    int64_t cudnn_rank = spatial_dims + 2;
+
     if (ShapeUtil::IsScalar(hlo.shape())) {
       Result result;
-      // cuDNN convolution tensors have a batch and a feature dimension in
-      // addition to spatial dimensions (at least 2 spatial dimensions for
-      // cuDNN).
-      int64_t spatial_dims =
-          std::max<int64_t>(2, dums_.input_spatial_dimensions_size());
-      result.sizes = std::vector<int64_t>(spatial_dims + 2, 1);
-      result.strides = std::vector<int64_t>(spatial_dims + 2, 1);
+      result.sizes = std::vector<int64_t>(cudnn_rank, 1);
+      result.strides = std::vector<int64_t>(cudnn_rank, 1);
       return result;
     }
-    if (hlo.shape().dimensions().size() == 1) {
+
+    int64_t conv_hlo_rank = dums_.input_spatial_dimensions_size() + 2;
+    if (hlo.shape().dimensions().size() < conv_hlo_rank) {
       Result result;
-      int64_t spatial_dims =
-          std::max<int64_t>(2, dums_.input_spatial_dimensions_size());
-      result.sizes = std::vector<int64_t>(spatial_dims + 2, 1);
-      result.strides = std::vector<int64_t>(spatial_dims + 2, 0);
-      result.sizes[1] = hlo.shape().dimensions(0);
-      result.strides[1] = 1;
+      result.sizes = std::vector<int64_t>(cudnn_rank, 1);
+      result.strides = std::vector<int64_t>(cudnn_rank, 0);
+
+      // If the parameter is consumed by a broadcast, map its dimensions to the
+      // corresponding cuDNN canonical axes (N, C, spatial...).
+      if (hlo.user_count() == 1 &&
+          hlo.users()[0]->opcode() == HloOpcode::kBroadcast) {
+        const auto& bcast_dims = hlo.users()[0]->dimensions();
+        for (int i = 0; i < bcast_dims.size(); ++i) {
+          int64_t cudnn_dim = HloDimToCudnnDim(bcast_dims[i]);
+          if (cudnn_dim >= 0 && cudnn_dim < cudnn_rank) {
+            result.sizes[cudnn_dim] = hlo.shape().dimensions(i);
+            result.strides[cudnn_dim] = 1;
+          }
+        }
+      } else if (hlo.shape().dimensions().size() == 1) {
+        // Fallback for un-broadcasted 1D parameters: assume channel bias [1, C,
+        // 1, 1].
+        result.sizes[1] = hlo.shape().dimensions(0);
+        result.strides[1] = 1;
+      } else {
+        return std::nullopt;
+      }
       return result;
     }
     // Placeholder FP32 data type here, it is not used.
diff --git a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler_test.cc b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler_test.cc
index eafc490c1b9e06..06293d680bd5f8 100644
--- a/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/cudnn_fusion_compiler_test.cc
@@ -207,6 +207,88 @@ TEST_F(CudnnFusionCompilerConstTest,
   EXPECT_GT(plan_count, 0);
 }
 
+TEST_F(CudnnFusionCompilerConstTest,
+       GetAvailablePlanCountFromConvolutionFusionWithBatchBroadcast) {
+  std::string hlo_text = R"(
+  fusion_batch_bcast {
+    p0 = f32[2,16,16,32] parameter(0)
+    p1 = f32[32,3,3,32] parameter(1)
+    conv = f32[2,16,16,32] convolution(p0, p1),
+      window={size=3x3 pad=1_1x1_1},
+      dim_labels=b01f_o01i->b01f,
+      convolution_kind=fprop
+    batch_vec = f32[2] parameter(2)
+    bcast = f32[2,16,16,32] broadcast(batch_vec), dimensions={0}
+    ROOT mul = f32[2,16,16,32] multiply(conv, bcast)
+  }
+
+  ENTRY e {
+    p0 = f32[2,16,16,32] parameter(0)
+    p1 = f32[32,3,3,32] parameter(1)
+    p2 = f32[2] parameter(2)
+    ROOT _ = f32[2,16,16,32] fusion(p0, p1, p2), kind=kCustom, calls=fusion_batch_bcast,
+      backend_config={
+        "fusion_backend_config": {
+          "kind": "__cudnn$$fusion",
+        }
+      }
+  })";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module,
+                       ParseAndReturnVerifiedModule(hlo_text));
+
+  const HloInstruction* root =
+      hlo_module->entry_computation()->root_instruction();
+  auto* fusion = Cast<HloFusionInstruction>(root);
+
+  ASSERT_OK_AND_ASSIGN(int plan_count,
+                       CuDnnFusionCompiler::GetAvailablePlanCount(
+                           stream_executor(),
+                           stream_executor()->GetDeviceDescription(), *fusion));
+  EXPECT_GT(plan_count, 0);
+}
+
+TEST_F(CudnnFusionCompilerConstTest,
+       GetAvailablePlanCountFrom1DConvolutionFusionWithBatchBroadcast) {
+  std::string hlo_text = R"(
+  fusion_1d_bcast {
+    p0 = f32[2,16,32] parameter(0)
+    p1 = f32[32,3,32] parameter(1)
+    conv = f32[2,16,32] convolution(p0, p1),
+      window={size=3 pad=1_1},
+      dim_labels=b0f_o0i->b0f,
+      convolution_kind=fprop
+    batch_vec = f32[2] parameter(2)
+    bcast = f32[2,16,32] broadcast(batch_vec), dimensions={0}
+    ROOT mul = f32[2,16,32] multiply(conv, bcast)
+  }
+
+  ENTRY e {
+    p0 = f32[2,16,32] parameter(0)
+    p1 = f32[32,3,32] parameter(1)
+    p2 = f32[2] parameter(2)
+    ROOT _ = f32[2,16,32] fusion(p0, p1, p2), kind=kCustom, calls=fusion_1d_bcast,
+      backend_config={
+        "fusion_backend_config": {
+          "kind": "__cudnn$$fusion",
+        }
+      }
+  })";
+
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module,
+                       ParseAndReturnVerifiedModule(hlo_text));
+
+  const HloInstruction* root =
+      hlo_module->entry_computation()->root_instruction();
+  auto* fusion = Cast<HloFusionInstruction>(root);
+
+  ASSERT_OK_AND_ASSIGN(int plan_count,
+                       CuDnnFusionCompiler::GetAvailablePlanCount(
+                           stream_executor(),
+                           stream_executor()->GetDeviceDescription(), *fusion));
+  EXPECT_GT(plan_count, 0);
+}
+
 }  // namespace
 }  // namespace gpu
 }  // namespace xla
diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
index 7715ce308eed56..a61accc39ba232 100644
--- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion.cc
@@ -1306,10 +1306,14 @@ FusionDecision ShouldFuseTranspose(const HloInstruction& transpose,
       dims->Indices(DotOperandDims::kContracting);
   absl::Span<const int64_t> non_contracting =
       dims->Indices(DotOperandDims::kNonContracting);
+  absl::Span<const int64_t> batch = dims->Indices(DotOperandDims::kBatch);
   if (!inverted_tracker->MapsToOneStride(contracting)) {
     return FusionDecision::Forbid(
         "Contracting dimension has non-contiguous section.");
   }
+  if (!inverted_tracker->MapsToOneStride(batch, /*allow_swaps=*/true)) {
+    return FusionDecision::Forbid("Batch dimension splits other dimensions.");
+  }
   if (operand_index == 1 &&
       !inverted_tracker->MapsToOneStride(non_contracting)) {
     return FusionDecision::Forbid(
diff --git a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc
index 3e9cd4b3d8144f..5113eefcd9eb41 100644
--- a/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc
+++ b/third_party/xla/xla/backends/gpu/transforms/gemm_fusion_test.cc
@@ -2394,6 +2394,62 @@ ENTRY main {
                   GmockMatch(TransposeOrBitcastTranspose())));
 }
 
+TEST_P(GemmFusionTestV2, AllowTransposeSwappingBatch) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
+HloModule module
+
+ENTRY main {
+  p_lhs = s8[8,128,64]{2,1,0} parameter(0)
+  cvt_lhs = bf16[8,128,64]{2,1,0} convert(p_lhs)
+  p_rhs = bf16[2,4,1,16,128]{4,3,2,1,0} parameter(1)
+
+  trans = bf16[4,2,1,16,128]{4,3,2,1,0} transpose(p_rhs), dimensions={1,0,2,3,4}
+  bitcast = bf16[8,16,128]{2,1,0} bitcast(trans)
+
+  ROOT dot = bf16[8,64,16]{2,1,0} dot(cvt_lhs, bitcast),
+    lhs_batch_dims={0}, lhs_contracting_dims={1},
+    rhs_batch_dims={0}, rhs_contracting_dims={2}
+}
+)"));
+
+  ASSERT_OK_AND_ASSIGN(bool changed,
+                       GemmFusion(gpu_version_).Run(module.get()));
+  EXPECT_TRUE(changed);
+  auto* fusion = module->entry_computation()->root_instruction();
+  EXPECT_THAT(fusion, GmockMatch(m::Fusion(m::Parameter(), m::Parameter())));
+}
+
+TEST_P(GemmFusionProfitabilityTest, DisallowTransposeInterleavingBatch) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
+                       ParseAndReturnVerifiedModule(R"(
+HloModule module
+
+ENTRY main {
+  p_lhs = s8[16,64,1024]{2,1,0} parameter(0)
+  cvt_lhs = bf16[16,64,1024]{2,1,0} convert(p_lhs)
+  p_rhs = bf16[4,16,4,64,32]{4,3,2,1,0} parameter(1)
+
+  trans = bf16[4,4,16,64,32]{4,3,2,1,0} transpose(p_rhs), dimensions={0,2,1,3,4}
+  bitcast = bf16[16,1024,32]{2,1,0} bitcast(trans)
+
+  ROOT dot = bf16[16,64,32]{2,1,0} dot(cvt_lhs, bitcast),
+    lhs_batch_dims={0}, lhs_contracting_dims={2},
+    rhs_batch_dims={0}, rhs_contracting_dims={1}
+}
+)"));
+
+  ASSERT_OK_AND_ASSIGN(bool changed,
+                       GemmFusion(gpu_version_).Run(module.get()));
+  EXPECT_TRUE(changed);
+  auto* fusion = module->entry_computation()->root_instruction();
+  EXPECT_THAT(fusion, GmockMatch(m::Fusion()));
+  EXPECT_THAT(fusion->operands(),
+              ::testing::UnorderedElementsAre(
+                  GmockMatch(m::Parameter()),
+                  GmockMatch(TransposeOrBitcastTranspose())));
+}
+
 TEST_P(GemmFusionTestV2, ConcatResetTrackerCrash) {
   ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnVerifiedModule(R"(
 HloModule m
diff --git a/third_party/xla/xla/hlo/analysis/alias_info.cc b/third_party/xla/xla/hlo/analysis/alias_info.cc
index b0d32815773303..cdf938757c834b 100644
--- a/third_party/xla/xla/hlo/analysis/alias_info.cc
+++ b/third_party/xla/xla/hlo/analysis/alias_info.cc
@@ -208,11 +208,10 @@ AliasInfo::GetInPlaceInputOutputPairs(const HloInstruction* user) const {
       return {};
     }
   }
-  if (user->opcode() == HloOpcode::kCustomCall) {
-    // Custom Calls previously assumed that aliased operands were
-    // forwarded, but now supports modification semantics.
-    const auto& aliasing_pairs =
-        Cast<HloCustomCallInstruction>(user)->output_to_operand_aliasing();
+  if (user->opcode() == HloOpcode::kCall ||
+      user->opcode() == HloOpcode::kCustomCall) {
+    const auto* callable = Cast<HloCallableInstruction>(user);
+    const auto& aliasing_pairs = callable->output_to_operand_aliasing();
     std::vector<std::pair<HloOperandIndex, ShapeIndex>> in_place_pairs;
     in_place_pairs.reserve(aliasing_pairs.size());
     for (const auto& pair : aliasing_pairs) {
diff --git a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc
index fee1d7580d3b89..1f87c30b7ba7d0 100644
--- a/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc
+++ b/third_party/xla/xla/hlo/analysis/hlo_alias_analysis_test.cc
@@ -24,8 +24,10 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "xla/hlo/analysis/alias_info.h"
 #include "xla/hlo/analysis/hlo_ordering.h"
+#include "xla/hlo/ir/hlo_casting_utils.h"
 #include "xla/hlo/ir/hlo_computation.h"
 #include "xla/hlo/ir/hlo_instruction.h"
+#include "xla/hlo/ir/hlo_instructions.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/hlo/testlib/hlo_hardware_independent_test_base.h"
 #include "xla/hlo/testlib/test.h"
@@ -1376,5 +1378,45 @@ ENTRY main {
             analysis.GetUniqueBufferAt(tc_operand1));
 }
 
+TEST_F(HloAliasAnalysisTest, CallOutputToOperandAliasing) {
+  absl::string_view hlo_string = R"(
+HloModule Module
+
+callme {
+  sub_p0 = f32[16] parameter(0)
+  sub_p1 = f32[16] parameter(1)
+  add0 = f32[16] add(sub_p0, sub_p1)
+  add1 = f32[16] add(sub_p0, sub_p0)
+  ROOT tuple = (f32[16], f32[16]) tuple(add0, add1)
+}
+
+ENTRY main {
+  entry_p0 = f32[16] parameter(0)
+  entry_p1 = f32[16] parameter(1)
+  ROOT call = (f32[16], f32[16]) call(entry_p0, entry_p1), to_apply=callme
+}
+)";
+  ASSERT_OK_AND_ASSIGN(module_, ParseAndReturnVerifiedModule(hlo_string));
+  auto* call = Cast<HloCallInstruction>(
+      module_->entry_computation()->GetInstructionWithName("call"));
+  call->set_output_to_operand_aliasing({{{0}, {0, {}}}, {{1}, {1, {}}}});
+
+  HloAliasAnalysis& analysis = RunAnalysis();
+
+  const HloInstruction* entry_p0 =
+      module_->entry_computation()->GetInstructionWithName("entry_p0");
+  const HloInstruction* entry_p1 =
+      module_->entry_computation()->GetInstructionWithName("entry_p1");
+
+  ASSERT_NE(entry_p0, nullptr);
+  ASSERT_NE(entry_p1, nullptr);
+  ASSERT_NE(call, nullptr);
+
+  EXPECT_EQ(analysis.GetUniqueBufferAt(entry_p0),
+            analysis.GetUniqueBufferAt(call, {0}));
+  EXPECT_EQ(analysis.GetUniqueBufferAt(entry_p1),
+            analysis.GetUniqueBufferAt(call, {1}));
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/analysis/shape_tracker.cc b/third_party/xla/xla/hlo/analysis/shape_tracker.cc
index a060b290381b76..e183e2d7a7e398 100644
--- a/third_party/xla/xla/hlo/analysis/shape_tracker.cc
+++ b/third_party/xla/xla/hlo/analysis/shape_tracker.cc
@@ -1255,7 +1255,8 @@ ShapeTracker::MapInputDimensionsToOutputUnordered(
   return kept_output_dims;
 }
 
-bool ShapeTracker::MapsToOneStride(absl::Span<const int64_t> input_dims) const {
+bool ShapeTracker::MapsToOneStride(absl::Span<const int64_t> input_dims,
+                                   bool allow_swaps) const {
   std::optional<std::vector<int64_t>> mapped =
       MapInputDimensionsToOutputUnordered(input_dims);
   if (!mapped.has_value()) {
@@ -1282,6 +1283,10 @@ bool ShapeTracker::MapsToOneStride(absl::Span<const int64_t> input_dims) const {
     return false;
   }
 
+  if (allow_swaps) {
+    return true;
+  }
+
   // Check for swaps.
   absl::StatusOr<ShapeTracker> narrowed = Narrow(input_dims);
   if (!narrowed.ok()) {
diff --git a/third_party/xla/xla/hlo/analysis/shape_tracker.h b/third_party/xla/xla/hlo/analysis/shape_tracker.h
index 2bf34cd3bbf039..f99794f4a0206f 100644
--- a/third_party/xla/xla/hlo/analysis/shape_tracker.h
+++ b/third_party/xla/xla/hlo/analysis/shape_tracker.h
@@ -143,9 +143,12 @@ class ShapeTracker {
       absl::Span<const int64_t> input_dims) const;
 
   // Returns true if the specified input dimensions map to a single contiguous
-  // stride in the output shape (i.e. they are not swapped, and there are no
-  // other non-degenerate dimensions between them in the output layout).
-  bool MapsToOneStride(absl::Span<const int64_t> input_dims) const;
+  // stride in the output shape. If `allow_swaps` is false (default), the
+  // dimensions must not be swapped with each other. In either case, there must
+  // be no other non-degenerate dimensions interleaved between them in the
+  // output layout.
+  bool MapsToOneStride(absl::Span<const int64_t> input_dims,
+                       bool allow_swaps = false) const;
 
   // Zips multiple ShapeTrackers into a single one.
   // For example, suppose we have two trackers:
diff --git a/third_party/xla/xla/hlo/analysis/shape_tracker_test.cc b/third_party/xla/xla/hlo/analysis/shape_tracker_test.cc
index 59febec3ccc9d7..bba5b2cfec9a62 100644
--- a/third_party/xla/xla/hlo/analysis/shape_tracker_test.cc
+++ b/third_party/xla/xla/hlo/analysis/shape_tracker_test.cc
@@ -1904,5 +1904,13 @@ TEST(ShapeTrackerMapsToOneStrideTest, ContiguousTransposeReshapeIsContiguous) {
   EXPECT_TRUE(tracker.MapsToOneStride({1, 2}));
 }
 
+TEST(ShapeTrackerMapsToOneStrideTest, AllowSwapsSwappedIsContiguous) {
+  Shape shape = ShapeUtil::MakeShape(F32, {2, 3, 4});
+  ShapeTracker tracker(shape);
+  ASSERT_TRUE(tracker.AppendTranspose({0, 2, 1}).ok());
+  EXPECT_FALSE(tracker.MapsToOneStride({1, 2}, /*allow_swaps=*/false));
+  EXPECT_TRUE(tracker.MapsToOneStride({1, 2}, /*allow_swaps=*/true));
+}
+
 }  // namespace
 }  // namespace xla
diff --git a/third_party/xla/xla/hlo/builder/lib/self_adjoint_eig_test.cc b/third_party/xla/xla/hlo/builder/lib/self_adjoint_eig_test.cc
index 5778da3b926576..467445ea097988 100644
--- a/third_party/xla/xla/hlo/builder/lib/self_adjoint_eig_test.cc
+++ b/third_party/xla/xla/hlo/builder/lib/self_adjoint_eig_test.cc
@@ -259,86 +259,6 @@ TEST_F(SelfAdjointEigTest, Test_Orthogonality_8x8) {
                              ErrorSpec(1e-3, 1e-3));
 }
 
-TEST_F(SelfAdjointEigTest, Test_Large_Magnitude_2x2) {
-  XlaBuilder builder(TestName());
-  float v = 1e20f;
-  Array2D<float> input{{v, v}, {v, v}};
-  std::vector<float> expected{0.0f, 2e20f};
-
-  XlaOp a;
-  auto a_data = CreateR2Parameter<float>(input, 0, "a", &builder, &a);
-  auto result = SelfAdjointEig(a, /*lower=*/true, /*max_iter=*/15,
-                               /*tol=*/1e-5, /*sort_eigenvalues=*/true);
-  Add(result.w, ZerosLike(result.w));
-
-  ComputeAndCompareR1<float>(&builder, expected, {&a_data},
-                             ErrorSpec(1e15f, 1e-4f));
-}
-
-TEST_F(SelfAdjointEigTest, Test_Large_Magnitude_3x3) {
-  XlaBuilder builder(TestName());
-  float v = 1e20f;
-  Array2D<float> input{{v, v, v}, {v, v, v}, {v, v, v}};
-  std::vector<float> expected{0.0f, 0.0f, 3e20f};
-
-  XlaOp a;
-  auto a_data = CreateR2Parameter<float>(input, 0, "a", &builder, &a);
-  auto result = SelfAdjointEig(a, /*lower=*/true, /*max_iter=*/15,
-                               /*tol=*/1e-5, /*sort_eigenvalues=*/true);
-  Add(result.w, ZerosLike(result.w));
-
-  ComputeAndCompareR1<float>(&builder, expected, {&a_data},
-                             ErrorSpec(1e15f, 1e-4f));
-}
-
-TEST_F(SelfAdjointEigTest, Test_Large_Magnitude_Complex_3x3) {
-  XlaBuilder builder(TestName());
-  float v = 1e20f;
-  Array<complex64> input = {
-      {complex64{v, 0.0f}, complex64{v, -v}, complex64{0.0f, 0.0f}},
-      {complex64{v, v}, complex64{v, 0.0f}, complex64{0.0f, 0.0f}},
-      {complex64{0.0f, 0.0f}, complex64{0.0f, 0.0f}, complex64{v, 0.0f}},
-  };
-  const Literal a_literal = LiteralUtil::CreateFromArray(input);
-  XlaOp a = Parameter(&builder, 0, a_literal.shape(), "a");
-  auto result = SelfAdjointEig(a);
-  ComputeMatmulVWVt(result, &builder);
-
-  ComputeAndCompareLiteral(&builder, LiteralUtil::CreateFromArray(input),
-                           {&a_literal}, ErrorSpec(1e15f, 1e-4f));
-}
-
-TEST_F(SelfAdjointEigTest, Test_Small_Magnitude_2x2) {
-  XlaBuilder builder(TestName());
-  float v = 1e-20f;
-  Array2D<float> input{{v, v}, {v, v}};
-  std::vector<float> expected{0.0f, 2e-20f};
-
-  XlaOp a;
-  auto a_data = CreateR2Parameter<float>(input, 0, "a", &builder, &a);
-  auto result = SelfAdjointEig(a, /*lower=*/true, /*max_iter=*/15,
-                               /*tol=*/1e-5, /*sort_eigenvalues=*/true);
-  Add(result.w, ZerosLike(result.w));
-
-  ComputeAndCompareR1<float>(&builder, expected, {&a_data},
-                             ErrorSpec(1e-25f, 1e-4f));
-}
-
-TEST_F(SelfAdjointEigTest, Test_Zero_Matrix_2x2) {
-  XlaBuilder builder(TestName());
-  Array2D<float> input{{0.0f, 0.0f}, {0.0f, 0.0f}};
-  std::vector<float> expected{0.0f, 0.0f};
-
-  XlaOp a;
-  auto a_data = CreateR2Parameter<float>(input, 0, "a", &builder, &a);
-  auto result = SelfAdjointEig(a, /*lower=*/true, /*max_iter=*/15,
-                               /*tol=*/1e-5, /*sort_eigenvalues=*/true);
-  Add(result.w, ZerosLike(result.w));
-
-  ComputeAndCompareR1<float>(&builder, expected, {&a_data},
-                             ErrorSpec(1e-6f, 1e-6f));
-}
-
 TEST_F(SelfAdjointEigTest, Wrong_Type_Int) {
   XlaBuilder builder(TestName());
 
diff --git a/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc b/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc
index 7400fb932f566f..3b6c0debf5a26e 100644
--- a/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc
+++ b/third_party/xla/xla/hlo/transforms/expanders/eigh_expander.cc
@@ -477,29 +477,6 @@ XlaOp EighExpander::BuildEigh(XlaOp a, bool lower, int64_t max_iter, float tol,
 
     a = Symmetrize(a, lower);
 
-    PrimitiveType real_type = primitive_util::IsComplexType(type)
-                                  ? primitive_util::ComplexComponentType(type)
-                                  : type;
-    XlaOp zero_real = Zero(builder, real_type);
-    XlaOp one_real = One(builder, real_type);
-    XlaOp abs_a = primitive_util::IsComplexType(type)
-                      ? Max(Abs(Real(a)), Abs(Imag(a)))
-                      : Abs(a);
-    XlaOp a_max =
-        Reduce(abs_a, zero_real, CreateScalarMaxComputation(real_type, builder),
-               {num_dims - 2, num_dims - 1});
-    XlaOp scale = Select(Eq(a_max, zero_real), one_real, a_max);
-
-    std::vector<int64_t> batch_broadcast_dims(num_batch_dims);
-    absl::c_iota(batch_broadcast_dims, 0);
-
-    XlaOp scale_a = primitive_util::IsComplexType(type)
-                        ? Complex(scale, ZerosLike(scale))
-                        : scale;
-    scale_a =
-        BroadcastInDim(scale_a, a_shape.dimensions(), batch_broadcast_dims);
-    a = a / scale_a;
-
     const int64_t k = CeilOfRatio(n, int64_t{2});
     // tl = A[:n // 2, :n // 2]
     // bl = A[n // 2:, :n // 2]
@@ -560,11 +537,6 @@ XlaOp EighExpander::BuildEigh(XlaOp a, bool lower, int64_t max_iter, float tol,
     }
     v = MaybeConjugate(TransposeInMinorDims(v), true);
 
-    ABSL_ASSIGN_OR_RETURN(Shape w_shape, builder->GetShape(w));
-    XlaOp scale_w =
-        BroadcastInDim(scale, w_shape.dimensions(), batch_broadcast_dims);
-    w = w * scale_w;
-
     if (sort_eigenvalues) {
       ABSL_RETURN_IF_ERROR(SortByEigenvalues(v, w));
     }
diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc
index 0b78e60c6d85a1..9e08a8071358bb 100644
--- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc
+++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.cc
@@ -2675,8 +2675,8 @@ PjRtCApiExecutable::GetOutputMemoryKinds() const {
   return std::vector<std::vector<absl::string_view>>{std::move(out)};
 }
 
-absl::StatusOr<std::vector<std::shared_ptr<HloModule>>>
-PjRtCApiExecutable::GetHloModules() const {
+absl::StatusOr<std::shared_ptr<HloModule>> PjRtCApiExecutable::GetHloModule()
+    const {
   auto* c_api = pjrt_c_api();
   auto* executable = c_executable();
   PJRT_Executable_OptimizedProgram_Args args;
@@ -2721,23 +2721,14 @@ PjRtCApiExecutable::GetHloModules() const {
     // equivalent) once implemented.
     mlir::MlirToHloConversionOptions options;
     options.return_tuple = false;
-    ABSL_ASSIGN_OR_RETURN(std::unique_ptr<xla::HloModule> hlo_module,
-                     mlir::ConvertMlirHloToHloModule(module.get(), options));
-
-    std::vector<std::shared_ptr<HloModule>> out;
-    out.push_back(std::move(hlo_module));
-    return out;
+    return mlir::ConvertMlirHloToHloModule(module.get(), options);
   }
 
   HloModuleProtoWithConfig proto;
   if (!proto.ParseFromString(code)) {
     return InvalidArgument("Failed to deserialize HloModuleProtoWithConfig");
   }
-  std::vector<std::shared_ptr<HloModule>> out;
-  ABSL_ASSIGN_OR_RETURN(std::unique_ptr<HloModule> module,
-                   HloModule::CreateFromProtoWithConfig(proto));
-  out.push_back(std::move(module));
-  return out;
+  return HloModule::CreateFromProtoWithConfig(proto);
 }
 
 absl::StatusOr<std::string> PjRtCApiExecutable::SerializeExecutable() const {
diff --git a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h
index 5c892ae2ecad28..e9fdfd162fc88f 100644
--- a/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h
+++ b/third_party/xla/xla/pjrt/c_api_client/pjrt_c_api_client.h
@@ -739,8 +739,7 @@ class PjRtCApiExecutable : public PjRtExecutable {
   absl::StatusOr<absl::flat_hash_map<std::string, PjRtValueType>>
   GetCostAnalysis() const override;
 
-  absl::StatusOr<std::vector<std::shared_ptr<HloModule>>> GetHloModules()
-      const override;
+  absl::StatusOr<std::shared_ptr<HloModule>> GetHloModule() const override;
 
   absl::StatusOr<CompiledMemoryStats> GetCompiledMemoryStats() const override {
     return pjrt::GetCompiledMemoryStats(c_api_, executable_.get());
diff --git a/third_party/xla/xla/pjrt/cpu/BUILD b/third_party/xla/xla/pjrt/cpu/BUILD
index 709693cc0a4c3a..e8170242698e7f 100644
--- a/third_party/xla/xla/pjrt/cpu/BUILD
+++ b/third_party/xla/xla/pjrt/cpu/BUILD
@@ -299,6 +299,9 @@ cc_library(
         "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/types:span",
         "@eigen_archive//:eigen3",
+        "@riegeli//riegeli/base:any",
+        "@riegeli//riegeli/bytes:reader",
+        "@riegeli//riegeli/messages:parse_message",
         "@tsl//tsl/platform:denormal",
         "@tsl//tsl/platform:fingerprint",
         "@tsl//tsl/platform:protobuf",
@@ -396,6 +399,8 @@ cc_library(
         "//xla/pjrt:pjrt_executable",
         "//xla/pjrt/plugin/xla_cpu:cpu_topology_description",
         "@com_google_absl//absl/status:statusor",
+        "@riegeli//riegeli/base:any",
+        "@riegeli//riegeli/bytes:reader",
     ],
     alwayslink = True,
 )
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.cc b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
index 64d780de8066ec..9fa31f3a93ae37 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.cc
@@ -44,6 +44,9 @@ limitations under the License.
 #include "absl/strings/string_view.h"
 #include "absl/synchronization/mutex.h"
 #include "absl/types/span.h"
+#include "riegeli/base/any.h"
+#include "riegeli/bytes/reader.h"
+#include "riegeli/messages/parse_message.h"
 #include "xla/array.h"
 #include "xla/backends/cpu/collectives/cpu_collectives.h"
 #include "xla/backends/cpu/constant_allocation.h"
@@ -565,15 +568,12 @@ PjRtCpuExecutable::GetOutputMemoryKinds() const {
   return out;
 }
 
-absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
-PjRtCpuClient::LoadSerializedExecutableInternal(
-    google::protobuf::io::ZeroCopyInputStream* stream,
-    std::optional<CompileOptions> options, const LoadOptions& load_options) {
+/*static*/ absl::StatusOr<std::unique_ptr<PjRtCpuExecutable>>
+PjRtCpuExecutable::Deserialize(riegeli::Any<riegeli::Reader*> reader,
+                               const xla::CpuTopologyDescription& topology,
+                               std::optional<CompileOptions>&& options) {
   ExecutableAndOptionsProto proto;
-  if (!proto.ParseFromZeroCopyStream(stream)) {
-    return Internal(
-        "PjRtCpuClient::DeserializeExecutable proto deserialization failed");
-  }
+  ABSL_RETURN_IF_ERROR(riegeli::ParseMessage(reader.get(), proto));
   CompileOptions compile_options;
   if (options.has_value()) {
     compile_options = *std::move(options);
@@ -599,10 +599,9 @@ PjRtCpuClient::LoadSerializedExecutableInternal(
   ABSL_RETURN_IF_ERROR(ParseDeviceAssignmentCompileOptions(
       compile_options.compile_portable_executable,
       &compile_options.executable_build_options,
-      [this](int num_replicas, int num_partitions) {
-        return topology().GetDefaultDeviceAssignment(process_index(),
-                                                     num_replicas, std::nullopt,
-                                                     num_partitions, nullptr);
+      [&topology](int num_replicas, int num_partitions) {
+        return topology.GetDefaultDeviceAssignment(
+            0, num_replicas, std::nullopt, num_partitions, nullptr);
       },
       &num_replicas, &num_partitions, &device_assignment));
 
@@ -630,37 +629,11 @@ PjRtCpuClient::LoadSerializedExecutableInternal(
                                : nullptr);
   }
 
-  auto cpu_executable = std::make_shared<PjRtCpuExecutable>(
+  auto cpu_executable = std::make_unique<PjRtCpuExecutable>(
       num_replicas, num_partitions, std::move(input_options),
       std::move(executable), std::move(result_buffer_indices), nullptr,
-      topology());
-  return LoadInternal(std::move(cpu_executable), std::move(device_assignment));
-}
-
-absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
-PjRtCpuClient::LoadSerializedExecutable(absl::string_view serialized,
-                                        std::optional<CompileOptions> options,
-                                        const LoadOptions& load_options) {
-  if (serialized.size() > std::numeric_limits<int>::max()) {
-    return Internal(
-        "PjRtCpuClient::DeserializeExecutable proto too large (>2GB)");
-  }
-  google::protobuf::io::ArrayInputStream stream(serialized.data(), serialized.size());
-  return LoadSerializedExecutableInternal(&stream, std::move(options),
-                                          load_options);
-}
-
-absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
-PjRtCpuClient::LoadSerializedExecutable(const absl::Cord& serialized,
-                                        std::optional<CompileOptions> options,
-                                        const LoadOptions& load_options) {
-  if (serialized.size() > std::numeric_limits<int>::max()) {
-    return Internal(
-        "PjRtCpuClient::DeserializeExecutable proto too large (>2GB)");
-  }
-  google::protobuf::io::CordInputStream stream(&serialized);
-  return LoadSerializedExecutableInternal(&stream, std::move(options),
-                                          load_options);
+      topology);
+  return cpu_executable;
 }
 
 absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
@@ -851,16 +824,6 @@ PjRtCpuRawClient::CompileAheadOfTime(const XlaComputation& computation,
                          &aot_options);
 }
 
-absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
-PjRtCpuClient::CompileAheadOfTimeAndLoad(
-    const XlaComputation& computation, CompileOptions options,
-    const AotCompilationOptions& aot_options) {
-  ABSL_ASSIGN_OR_RETURN(auto executable, raw_client()->CompileAheadOfTime(
-                                        computation, options, topology(),
-                                        process_index(), aot_options));
-  return Load(std::move(executable), LoadOptions());
-}
-
 struct CpuCompilationParams {
   LayoutCanonicalizationCallback layout_canonicalization_callback = nullptr;
   std::optional<int> num_threads = std::nullopt;
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_client.h b/third_party/xla/xla/pjrt/cpu/cpu_client.h
index 2276bb60252b63..01ca97aafd81d3 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_client.h
+++ b/third_party/xla/xla/pjrt/cpu/cpu_client.h
@@ -267,27 +267,6 @@ class PjRtCpuClient final : public CommonPjRtClientImpl {
       std::shared_ptr<PjRtExecutable> executable,
       const LoadOptions& load_options) override;
 
-  // TODO(b/403584258): PJRT wants to have just one simple Compile API. When the
-  // CPU runtime stops supporting the legacy runtime we will unify our compile
-  // paths better and this will be redundant.
-  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
-  CompileAheadOfTimeAndLoad(const XlaComputation& computation,
-                            CompileOptions options,
-                            const AotCompilationOptions& aot_options);
-
-  // For PjRtCpuClient, `options` is mandatory.
-  // This function returns an InvalidArgument error if `std::nullopt` is passed.
-  // TODO(b/237720161): make it actually optional
-  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
-  LoadSerializedExecutable(absl::string_view serialized,
-                           std::optional<CompileOptions> options,
-                           const LoadOptions& load_options) override;
-
-  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
-  LoadSerializedExecutable(const absl::Cord& serialized,
-                           std::optional<CompileOptions> options,
-                           const LoadOptions& load_options) override;
-
   bool IsOnCpu(PjRtMemorySpace* memory_space) override { return true; }
 
   const xla::CpuTopologyDescription& topology() const {
@@ -319,11 +298,6 @@ class PjRtCpuClient final : public CommonPjRtClientImpl {
   absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>> LoadInternal(
       std::shared_ptr<PjRtCpuExecutable> cpu_executable,
       std::shared_ptr<DeviceAssignment> device_assignment);
-
-  absl::StatusOr<std::unique_ptr<PjRtLoadedExecutable>>
-  LoadSerializedExecutableInternal(google::protobuf::io::ZeroCopyInputStream* stream,
-                                   std::optional<CompileOptions> options,
-                                   const LoadOptions& load_options);
 };
 
 class PjRtCpuLoadedExecutable;
@@ -399,10 +373,8 @@ class PjRtCpuExecutable final : public PjRtExecutable {
     return cpu_executable_->SizeOfGeneratedCodeInBytes();
   }
 
-  absl::StatusOr<std::vector<std::shared_ptr<HloModule>>> GetHloModules()
-      const override {
-    return std::vector<std::shared_ptr<HloModule>>{
-        cpu_executable_->shared_module()};
+  absl::StatusOr<std::shared_ptr<HloModule>> GetHloModule() const override {
+    return cpu_executable_->shared_module();
   }
 
   absl::StatusOr<std::vector<std::vector<absl::string_view>>>
@@ -431,6 +403,11 @@ class PjRtCpuExecutable final : public PjRtExecutable {
 
   const CompileOptions& compile_options() const { return compile_options_; }
 
+  static absl::StatusOr<std::unique_ptr<PjRtCpuExecutable>> Deserialize(
+      riegeli::Any<riegeli::Reader*> reader,
+      const xla::CpuTopologyDescription& topology,
+      std::optional<CompileOptions>&& options);
+
  private:
   friend class PjRtCpuClient;
   friend class CpuPjRtRawLoadedExecutable;
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.cc b/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.cc
index 910dc9460103fc..39c600b97c3dcb 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.cc
+++ b/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.cc
@@ -16,6 +16,7 @@ limitations under the License.
 #include "xla/pjrt/cpu/cpu_pjrt_compiler.h"
 
 #include <memory>
+#include <optional>
 #include <string>
 #include <utility>
 
@@ -25,6 +26,8 @@ limitations under the License.
 #include "absl/status/status_macros.h"
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
+#include "riegeli/base/any.h"
+#include "riegeli/bytes/reader.h"
 #include "xla/hlo/builder/xla_computation.h"
 #include "xla/pjrt/cpu/cpu_client.h"
 #include "xla/pjrt/maybe_owning_mlir_module.h"
@@ -99,6 +102,16 @@ CpuPjRtCompiler::DeserializePjRtTopologyDescription(
   return CpuTopologyDescription::FromProto(proto);
 }
 
+absl::StatusOr<std::unique_ptr<PjRtExecutable>>
+CpuPjRtCompiler::DeserializeExecutable(
+    const PjRtTopologyDescription& topology,
+    riegeli::Any<riegeli::Reader*> reader,
+    std::optional<CompileOptions>&& options) {
+  ABSL_ASSIGN_OR_RETURN(auto* cpu_topology, GetCpuTopology(topology));
+  return PjRtCpuExecutable::Deserialize(std::move(reader), *cpu_topology,
+                                        std::move(options));
+}
+
 }  // namespace xla::cpu
 
 // Doesn't use STREAM_EXECUTOR_REGISTER_MODULE_INITIALIZER to ensure it is
diff --git a/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.h b/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.h
index 28672b3fd0876f..6cb1307d82619c 100644
--- a/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.h
+++ b/third_party/xla/xla/pjrt/cpu/cpu_pjrt_compiler.h
@@ -17,9 +17,12 @@ limitations under the License.
 #define XLA_PJRT_CPU_CPU_PJRT_COMPILER_H_
 
 #include <memory>
+#include <optional>
 #include <string>
 
 #include "absl/status/statusor.h"
+#include "riegeli/base/any.h"
+#include "riegeli/bytes/reader.h"
 #include "xla/hlo/builder/xla_computation.h"
 #include "xla/pjrt/maybe_owning_mlir_module.h"
 #include "xla/pjrt/pjrt_compiler.h"
@@ -40,6 +43,12 @@ class CpuPjRtCompiler : public PjRtCompiler {
       CompileOptions options, MaybeOwningMlirModule module,
       const PjRtTopologyDescription& topology, PjRtClient* client) override;
 
+  // Deserializes a serialized executable.
+  absl::StatusOr<std::unique_ptr<PjRtExecutable>> DeserializeExecutable(
+      const PjRtTopologyDescription& topology,
+      riegeli::Any<riegeli::Reader*> reader,
+      std::optional<CompileOptions>&& options) override;
+
   // Deserializes a PjRtTopologyDescription from a string.
   absl::StatusOr<std::unique_ptr<PjRtTopologyDescription>>
   DeserializePjRtTopologyDescription(
diff --git a/third_party/xla/xla/pjrt/gpu/BUILD b/third_party/xla/xla/pjrt/gpu/BUILD
index ec0db00389c5fa..4b331c67bc86ca 100644
--- a/third_party/xla/xla/pjrt/gpu/BUILD
+++ b/third_party/xla/xla/pjrt/gpu/BUILD
@@ -917,6 +917,7 @@ xla_test(
         "//xla/service:compiled_module",
         "//xla/service:compiler",
         "//xla/service:gpu_topology",
+        "//xla/service:mock_compiled_module",
         "//xla/service:mock_compiler",
         "//xla/service:platform_util",
         "//xla/stream_executor:stream_executor_h",
diff --git a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc
index bd28695ad9d68a..cdb2f122b4b003 100644
--- a/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc
+++ b/third_party/xla/xla/pjrt/gpu/se_gpu_pjrt_compiler_test.cc
@@ -55,6 +55,7 @@ limitations under the License.
 #include "xla/service/compiled_module.h"
 #include "xla/service/compiler.h"
 #include "xla/service/gpu_topology.h"
+#include "xla/service/mock_compiled_module.h"
 #include "xla/service/mock_compiler.h"
 #include "xla/service/platform_util.h"
 #include "xla/shape_layout.h"
@@ -67,7 +68,8 @@ namespace xla {
 namespace {
 using ::absl_testing::IsOkAndHolds;
 using ::testing::_;
-using ::testing::IsEmpty;
+using ::testing::ByMove;
+using ::testing::ElementsAre;
 using ::testing::IsNull;
 using ::testing::Optional;
 using ::testing::Property;
@@ -313,6 +315,17 @@ TEST(StreamExecutorGpuCompilerTest, GetTargetRuntimeAbiVersion) {
   EXPECT_OK(runtime_abi_version->IsCompatibleWith(*executable_abi_version));
 }
 
+std::vector<std::unique_ptr<CompiledModule>> MakeMockCompiledModule(
+    std::shared_ptr<HloModule> hlo_module =
+        std::make_shared<HloModule>("name", HloModuleConfig())) {
+  auto mock_compiled_module = std::make_unique<MockCompiledModule>();
+  EXPECT_CALL(*mock_compiled_module, shared_optimized_module())
+      .WillRepeatedly(Return(hlo_module));
+  std::vector<std::unique_ptr<CompiledModule>> aot_results;
+  aot_results.push_back(std::move(mock_compiled_module));
+  return aot_results;
+}
+
 TEST(StreamExecutorGpuCompilerTest, DevicelessCompilation) {
   auto mock_compiler = std::make_unique<MockCompiler>();
   MockCompiler& mock_compiler_ref = *mock_compiler;
@@ -329,6 +342,8 @@ TEST(StreamExecutorGpuCompilerTest, DevicelessCompilation) {
   StreamExecutorGpuTopologyDescription topology_description(
       CudaId(), CudaName(), gpu_topology);
 
+  auto hlo_module = std::make_shared<HloModule>("name", HloModuleConfig());
+
   // We expect that the underlying compiler is called with no executor given.
   EXPECT_CALL(mock_compiler_ref, PlatformId)
       .WillRepeatedly(Return(stream_executor::cuda::kCudaPlatformId));
@@ -336,14 +351,15 @@ TEST(StreamExecutorGpuCompilerTest, DevicelessCompilation) {
   EXPECT_CALL(mock_compiler_ref,
               CompileAheadOfTime(
                   _, Property(&AotCompilationOptions::executor, IsNull())))
-      .WillOnce(Return(std::vector<std::unique_ptr<CompiledModule>>{}));
+      .WillOnce(Return(ByMove(MakeMockCompiledModule(hlo_module))));
 
   ASSERT_OK_AND_ASSIGN(XlaComputation computation, GetXlaComputation(kProgram));
   ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<PjRtExecutable> executable,
       pjrt_compiler.Compile(CompileOptions(), computation, topology_description,
                             /*client=*/nullptr));
-  EXPECT_THAT(executable->GetHloModules(), IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(executable->GetHloModules(),
+              IsOkAndHolds(ElementsAre(hlo_module)));
 }
 
 TEST(StreamExecutorGpuCompilerTest, CrossCompilation) {
@@ -372,6 +388,8 @@ TEST(StreamExecutorGpuCompilerTest, CrossCompilation) {
   se::StreamExecutor* stream_executor =
       se_gpu_client->client()->backend().default_stream_executor();
 
+  auto hlo_module = std::make_shared<HloModule>("name", HloModuleConfig());
+
   // We expect that the underlying compiler is called with the executor from the
   // PjRt client.
   EXPECT_CALL(mock_compiler_ref, PlatformId)
@@ -390,14 +408,15 @@ TEST(StreamExecutorGpuCompilerTest, CrossCompilation) {
                            Property(&GpuTopology::num_partitions, 1),
                            Property(&GpuTopology::num_hosts_per_partition, 1),
                            Property(&GpuTopology::num_devices_per_host, 1)))))))
-      .WillOnce(Return(std::vector<std::unique_ptr<CompiledModule>>{}));
+      .WillOnce(Return(ByMove(MakeMockCompiledModule(hlo_module))));
 
   ASSERT_OK_AND_ASSIGN(XlaComputation computation, GetXlaComputation(kProgram));
   ASSERT_OK_AND_ASSIGN(
       std::unique_ptr<PjRtExecutable> executable,
       pjrt_compiler.Compile(CompileOptions(), computation, topology_description,
                             client.get()));
-  EXPECT_THAT(executable->GetHloModules(), IsOkAndHolds(IsEmpty()));
+  EXPECT_THAT(executable->GetHloModules(),
+              IsOkAndHolds(ElementsAre(hlo_module)));
 }
 
 absl::StatusOr<std::shared_ptr<GpuTopology>> GetSampleH100basedGpuTopology() {
@@ -436,12 +455,12 @@ TEST(StreamExecutorGpuCompilerTest, AutoLayoutIsPropagatedInCrossCompilation) {
       .WillRepeatedly(Return(stream_executor::cuda::kCudaPlatformId));
   EXPECT_CALL(mock_compiler_ref, Compile).Times(0);
 
-  std::unique_ptr<HloModule> hlo_module;
+  std::shared_ptr<HloModule> hlo_module;
   EXPECT_CALL(mock_compiler_ref, CompileAheadOfTime)
       .WillOnce([&](std::unique_ptr<HloModule> module,
                     const AotCompilationOptions& options) {
         hlo_module = std::move(module);
-        return std::vector<std::unique_ptr<CompiledModule>>{};
+        return MakeMockCompiledModule(hlo_module);
       });
 
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<PjRtClient> client,
@@ -475,12 +494,12 @@ TEST(StreamExecutorGpuCompilerTest,
       .WillRepeatedly(Return(stream_executor::cuda::kCudaPlatformId));
   EXPECT_CALL(mock_compiler_ref, Compile).Times(0);
 
-  std::unique_ptr<HloModule> hlo_module;
+  std::shared_ptr<HloModule> hlo_module;
   EXPECT_CALL(mock_compiler_ref, CompileAheadOfTime)
       .WillOnce([&](std::unique_ptr<HloModule> module,
                     const AotCompilationOptions& options) {
         hlo_module = std::move(module);
-        return std::vector<std::unique_ptr<CompiledModule>>{};
+        return MakeMockCompiledModule(hlo_module);
       });
 
   ASSERT_OK_AND_ASSIGN(
@@ -512,12 +531,12 @@ TEST(StreamExecutorGpuCompilerTest,
       .WillRepeatedly(Return(stream_executor::cuda::kCudaPlatformId));
   EXPECT_CALL(mock_compiler_ref, Compile).Times(0);
 
-  std::unique_ptr<HloModule> hlo_module;
+  std::shared_ptr<HloModule> hlo_module;
   EXPECT_CALL(mock_compiler_ref, CompileAheadOfTime)
       .WillOnce([&](std::unique_ptr<HloModule> module,
                     const AotCompilationOptions& options) {
         hlo_module = std::move(module);
-        return std::vector<std::unique_ptr<CompiledModule>>{};
+        return MakeMockCompiledModule(hlo_module);
       });
 
   CompileOptions options{};
diff --git a/third_party/xla/xla/pjrt/pjrt_executable.cc b/third_party/xla/xla/pjrt/pjrt_executable.cc
index 50dd2df937407f..80708b7d4614fd 100644
--- a/third_party/xla/xla/pjrt/pjrt_executable.cc
+++ b/third_party/xla/xla/pjrt/pjrt_executable.cc
@@ -315,6 +315,12 @@ void GetOpSharding(std::vector<OpSharding>& out, const OpSharding& sharding) {
   }
 }
 
+absl::StatusOr<std::vector<std::shared_ptr<HloModule>>>
+PjRtExecutable::GetHloModules() const {
+  ABSL_ASSIGN_OR_RETURN(std::shared_ptr<HloModule> hlo_module, GetHloModule());
+  return std::vector<std::shared_ptr<HloModule>>{std::move(hlo_module)};
+}
+
 std::optional<std::vector<OpSharding>> PjRtExecutable::GetOutputShardings()
     const {
   auto modules = GetHloModules();
diff --git a/third_party/xla/xla/pjrt/pjrt_executable.h b/third_party/xla/xla/pjrt/pjrt_executable.h
index ee23b502ae4aa3..4a7c456c8457a1 100644
--- a/third_party/xla/xla/pjrt/pjrt_executable.h
+++ b/third_party/xla/xla/pjrt/pjrt_executable.h
@@ -367,9 +367,14 @@ class PjRtExecutable {
   // Unique name for this executable, e.g., HloModule name.
   virtual absl::string_view name() const = 0;
 
+  // Return HloModule (optimized).
+  virtual absl::StatusOr<std::shared_ptr<HloModule>> GetHloModule() const {
+    return absl::UnimplementedError("GetHloModule is not implemented");
+  }
+
   // Return an array of HloModule (optimized) per partition.
   virtual absl::StatusOr<std::vector<std::shared_ptr<HloModule>>>
-  GetHloModules() const = 0;
+  GetHloModules() const;
 
   // Returns an output Shape per program, the size should be equal to
   // `GetHloModules()`.
diff --git a/third_party/xla/xla/pjrt/se/BUILD b/third_party/xla/xla/pjrt/se/BUILD
index aaeb4658eabdde..154374dbfde4e2 100644
--- a/third_party/xla/xla/pjrt/se/BUILD
+++ b/third_party/xla/xla/pjrt/se/BUILD
@@ -178,11 +178,13 @@ xla_cc_test(
     srcs = ["stream_executor_executable_test.cc"],
     deps = [
         ":stream_executor_executable",
+        "//xla/hlo/ir:hlo",
         "//xla/pjrt:pjrt_abi_version",
         "//xla/pjrt:pjrt_common",
         "//xla/pjrt:pjrt_executable",
         "//xla/pjrt/proto:pjrt_abi_version_proto_cc",
         "//xla/service:compiled_module",
+        "//xla/service:hlo_module_config",
         "//xla/service:mock_compiled_module",
         "//xla/stream_executor/abi:executable_abi_version",
         "//xla/stream_executor/abi:executable_abi_version_proto_cc",
diff --git a/third_party/xla/xla/pjrt/se/stream_executor_executable.cc b/third_party/xla/xla/pjrt/se/stream_executor_executable.cc
index 61e578cda19a00..94eaddff745c85 100644
--- a/third_party/xla/xla/pjrt/se/stream_executor_executable.cc
+++ b/third_party/xla/xla/pjrt/se/stream_executor_executable.cc
@@ -169,6 +169,7 @@ StreamExecutorExecutable::StreamExecutorExecutable(
   if (mod != nullptr) {
     hlo_module_ = mod->shared_optimized_module();
   }
+  CHECK(hlo_module_ != nullptr);
 }
 
 StreamExecutorExecutable::StreamExecutorExecutable(
@@ -193,6 +194,7 @@ StreamExecutorExecutable::StreamExecutorExecutable(
   if (local_exec != nullptr) {
     hlo_module_ = local_exec->executable()->shared_module();
   }
+  CHECK(hlo_module_ != nullptr);
 }
 
 absl::StatusOr<CompiledMemoryStats>
diff --git a/third_party/xla/xla/pjrt/se/stream_executor_executable.h b/third_party/xla/xla/pjrt/se/stream_executor_executable.h
index 31b14369015a11..cb0175a729b5d2 100644
--- a/third_party/xla/xla/pjrt/se/stream_executor_executable.h
+++ b/third_party/xla/xla/pjrt/se/stream_executor_executable.h
@@ -78,12 +78,8 @@ class StreamExecutorExecutable : public PjRtExecutable {
   absl::StatusOr<CompileOptions> GetCompileOptions() const override {
     return compile_options_;
   }
-  absl::StatusOr<std::vector<std::shared_ptr<HloModule>>> GetHloModules()
-      const override {
-    if (hlo_module_ == nullptr) {
-      return std::vector<std::shared_ptr<HloModule>>{};
-    }
-    return std::vector<std::shared_ptr<HloModule>>{hlo_module_};
+  absl::StatusOr<std::shared_ptr<HloModule>> GetHloModule() const override {
+    return hlo_module_;
   }
 
   const std::shared_ptr<HloModule>& hlo_module() const { return hlo_module_; }
diff --git a/third_party/xla/xla/pjrt/se/stream_executor_executable_test.cc b/third_party/xla/xla/pjrt/se/stream_executor_executable_test.cc
index c22e5e255d7bf9..2da6cb9a4bab5b 100644
--- a/third_party/xla/xla/pjrt/se/stream_executor_executable_test.cc
+++ b/third_party/xla/xla/pjrt/se/stream_executor_executable_test.cc
@@ -21,11 +21,13 @@ limitations under the License.
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "xla/hlo/ir/hlo_module.h"
 #include "xla/pjrt/pjrt_abi_version.h"
 #include "xla/pjrt/pjrt_common.h"
 #include "xla/pjrt/pjrt_executable.h"
 #include "xla/pjrt/proto/pjrt_abi_version.pb.h"
 #include "xla/service/compiled_module.h"
+#include "xla/service/hlo_module_config.h"
 #include "xla/service/mock_compiled_module.h"
 #include "xla/stream_executor/abi/executable_abi_version.h"
 #include "xla/stream_executor/abi/executable_abi_version.pb.h"
@@ -49,6 +51,9 @@ TEST(StreamExecutorExecutableTest, GetAbiVersion) {
   constexpr PjRtPlatformId kPlatformId = 42;
 
   auto module = std::make_unique<MockCompiledModule>();
+  auto hlo_module = std::make_shared<HloModule>("name", HloModuleConfig());
+  EXPECT_CALL(*module, shared_optimized_module())
+      .WillRepeatedly(Return(hlo_module));
   EXPECT_CALL(*module, GetExecutableAbiVersion())
       .WillOnce(Return(executable_abi_version));
   StreamExecutorExecutable executable(kPlatformId, CompileOptions(),
diff --git a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable_test_lib.cc b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable_test_lib.cc
index 07be8793b322fe..7730983c23967a 100644
--- a/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable_test_lib.cc
+++ b/third_party/xla/xla/python/ifrt/ir/ifrt_ir_loaded_executable_test_lib.cc
@@ -1473,13 +1473,19 @@ module {
                                   ExecuteOptionsWithFillStatus(), devices));
   EXPECT_THAT(
       result.status.Await(),
-      StatusIs(absl::StatusCode::kInvalidArgument,
-               HasSubstr("Invalid buffer passed to Execute() as argument 0")));
+      AnyOf(StatusIs(
+                absl::StatusCode::kInvalidArgument,
+                HasSubstr("Invalid buffer passed to Execute() as argument 0")),
+            StatusIs(absl::StatusCode::kInvalidArgument,
+                     HasSubstr("donates or deletes twice an array"))));
   ASSERT_EQ(result.outputs.size(), 1);
   EXPECT_THAT(
       result.outputs[0]->GetReadyFuture().Await(),
-      StatusIs(absl::StatusCode::kInvalidArgument,
-               HasSubstr("Invalid buffer passed to Execute() as argument 0")));
+      AnyOf(StatusIs(
+                absl::StatusCode::kInvalidArgument,
+                HasSubstr("Invalid buffer passed to Execute() as argument 0")),
+            StatusIs(absl::StatusCode::kInvalidArgument,
+                     HasSubstr("donates or deletes twice an array"))));
 }
 
 TEST_F(IfrtIrLoadedExecutableTest, DonatingTwiceAliasedBufferThrowsError) {
@@ -1533,6 +1539,8 @@ module {
           StatusIs(absl::StatusCode::kInvalidArgument,
                    HasSubstr(
                        "Attempt to donate the same buffer twice in Execute()")),
+          StatusIs(absl::StatusCode::kInvalidArgument,
+                   HasSubstr("donates or deletes twice an array")),
           StatusIs(
               absl::StatusCode::kNotFound,
               HasSubstr(
@@ -1544,6 +1552,8 @@ module {
           StatusIs(absl::StatusCode::kInvalidArgument,
                    HasSubstr(
                        "Attempt to donate the same buffer twice in Execute()")),
+          StatusIs(absl::StatusCode::kInvalidArgument,
+                   HasSubstr("donates or deletes twice an array")),
           StatusIs(
               absl::StatusCode::kNotFound,
               HasSubstr(
diff --git a/third_party/xla/xla/service/gpu/BUILD b/third_party/xla/xla/service/gpu/BUILD
index 1357ef9030f976..eb3e044fe4fc25 100644
--- a/third_party/xla/xla/service/gpu/BUILD
+++ b/third_party/xla/xla/service/gpu/BUILD
@@ -1971,6 +1971,7 @@ cc_library(
         "//xla/hlo/transforms/simplifiers:broadcast_canonicalizer",
         "//xla/hlo/transforms/simplifiers:conditional_canonicalizer",
         "//xla/hlo/transforms/simplifiers:convert_mover",
+        "//xla/hlo/transforms/simplifiers:degenerate_dimension_rewriter",
         "//xla/hlo/transforms/simplifiers:dot_merger",
         "//xla/hlo/transforms/simplifiers:dynamic_dimension_simplifier",
         "//xla/hlo/transforms/simplifiers:flatten_call_graph",
diff --git a/third_party/xla/xla/service/gpu/gpu_compiler.cc b/third_party/xla/xla/service/gpu/gpu_compiler.cc
index 97d916a9feb43f..4ca8fef163aebd 100644
--- a/third_party/xla/xla/service/gpu/gpu_compiler.cc
+++ b/third_party/xla/xla/service/gpu/gpu_compiler.cc
@@ -209,6 +209,7 @@ limitations under the License.
 #include "xla/hlo/transforms/simplifiers/broadcast_canonicalizer.h"
 #include "xla/hlo/transforms/simplifiers/conditional_canonicalizer.h"
 #include "xla/hlo/transforms/simplifiers/convert_mover.h"
+#include "xla/hlo/transforms/simplifiers/degenerate_dimension_rewriter.h"
 #include "xla/hlo/transforms/simplifiers/dot_merger.h"
 #include "xla/hlo/transforms/simplifiers/dynamic_dimension_simplifier.h"
 #include "xla/hlo/transforms/simplifiers/flatten_call_graph.h"
@@ -977,6 +978,14 @@ absl::Status RunOptimizationPasses(
     pipeline.AddPass<ScatterSliceSimplifier>();
     pipeline.AddPass<DotStrengthReduction>(
         gpu_target_config.device_description.gpu_compute_capability());
+
+    // It's important to run AlgebraicSimplifier after
+    // DegenerateDimensionRewriter before ReshapeMover.
+    // DegenerateDimensionRewriter introduces reshape to remove size-1 dims from
+    // ops like iota and broadcast, and algebraic simplifier has patterns to
+    // fold reshape(iota) and reshape(broadcast). If we run ReshapeMover first,
+    // it will move these reshapes down the graph, and prevent the folding.
+    pipeline.AddPass<DegenerateDimensionRewriter>();
     pipeline.AddPass<GpuAlgebraicSimplifier>(layout_insensitive_algsimp_opts,
                                              gpu_version);
     pipeline.AddPass<SortSimplifier>();
diff --git a/third_party/xla/xla/service/hlo_verifier.cc b/third_party/xla/xla/service/hlo_verifier.cc
index 2ef568d5658deb..b6ab4b9cf7a4b5 100644
--- a/third_party/xla/xla/service/hlo_verifier.cc
+++ b/third_party/xla/xla/service/hlo_verifier.cc
@@ -4387,17 +4387,35 @@ absl::Status InstructionVerifier::HandleWhile(HloInstruction* xla_while) {
 
 absl::Status InstructionVerifier::HandleCall(HloInstruction* call) {
   if (opts_.verify_call_nested_computation_thread_name) {
-    return CheckCallableInstructionThreadName(call);
+    ABSL_RETURN_IF_ERROR(CheckCallableInstructionThreadName(call));
+  }
+  const auto* callable = Cast<const HloCallableInstruction>(call);
+  for (const auto& pair : callable->output_to_operand_aliasing()) {
+    TF_RET_CHECK(pair.second.first < callable->operand_count())
+        << "Invalid aliasing operand index.";
+    TF_RET_CHECK(ShapeUtil::IndexIsValid(
+        callable->operand(pair.second.first)->shape(), pair.second.second))
+        << "Invalid aliasing operand shape index.";
+    TF_RET_CHECK(ShapeUtil::IndexIsValid(callable->shape(), pair.first))
+        << "Invalid aliasing output shape index.";
+    const Shape& output_subshape =
+        ShapeUtil::GetSubshape(callable->shape(), pair.first);
+    const Shape& operand_subshape = ShapeUtil::GetSubshape(
+        callable->operand(pair.second.first)->shape(), pair.second.second);
+    if (opts_.layout_sensitive) {
+      TF_RET_CHECK(Shape::Equal().IgnoreDynamicDimension()(operand_subshape,
+                                                           output_subshape))
+          << "Different aliasing shapes: "
+          << operand_subshape.ToString(/*print_layout=*/true) << " vs "
+          << output_subshape.ToString(/*print_layout=*/true);
+    } else {
+      TF_RET_CHECK(Shape::Equal().IgnoreDynamicDimension().IgnoreLayout()(
+          operand_subshape, output_subshape))
+          << "Different aliasing shapes: "
+          << operand_subshape.ToString(/*print_layout=*/false) << " vs "
+          << output_subshape.ToString(/*print_layout=*/false);
+    }
   }
-
-  // As opposed to other callable instructions, nothing respects input/output
-  // aliasing for call instructions, so make sure it's not set.
-  const HloCallableInstruction* callable =
-      DynCast<const HloCallableInstruction>(call);
-  TF_RET_CHECK(callable != nullptr);
-  TF_RET_CHECK(callable->output_to_operand_aliasing().empty())
-      << "Call instruction " << call->ToString()
-      << " may not have an output-to-operand aliasing set.";
   return absl::OkStatus();
 }
 
diff --git a/third_party/xla/xla/service/hlo_verifier_test.cc b/third_party/xla/xla/service/hlo_verifier_test.cc
index 26f5b36fc21d51..117e53583e19bb 100644
--- a/third_party/xla/xla/service/hlo_verifier_test.cc
+++ b/third_party/xla/xla/service/hlo_verifier_test.cc
@@ -326,10 +326,133 @@ TEST_F(HloVerifierTest, CheckCallOperandOutputAliasing) {
       module->entry_computation()->GetInstructionWithName("mycall"))
       ->set_output_to_operand_aliasing({{{}, {0, {}}}});
 
+  EXPECT_OK(verifier().Run(module.get()).status());
+}
+
+TEST_F(HloVerifierTest, CheckCallTupleOperandOutputAliasing) {
+  constexpr absl::string_view hlo = R"(
+    HloModule Module
+
+    callme {
+      p0 = (s32[], f32[4]) parameter(0)
+      p1 = f32[4] parameter(1)
+      p0_0 = s32[] get-tuple-element(p0), index=0
+      ROOT result = (s32[], f32[4]) tuple(p0_0, p1)
+    }
+
+    ENTRY entry {
+      p0 = (s32[], f32[4]) parameter(0)
+      p1 = f32[4] parameter(1)
+      ROOT mycall = (s32[], f32[4]) call(p0, p1), to_apply=callme
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo));
+
+  Cast<HloCallInstruction>(
+      module->entry_computation()->GetInstructionWithName("mycall"))
+      ->set_output_to_operand_aliasing({{{0}, {0, {0}}}, {{1}, {1, {}}}});
+
+  EXPECT_OK(verifier().Run(module.get()).status());
+}
+
+TEST_F(HloVerifierTest, CheckCallOperandOutputAliasingInvalidOperandIndex) {
+  constexpr absl::string_view hlo = R"(
+    HloModule Module
+
+    callme {
+      ROOT param = (s32[], f32[4]) parameter(0)
+    }
+
+    ENTRY entry {
+      p0 = (s32[], f32[4]) parameter(0)
+      ROOT mycall = (s32[], f32[4]) call(p0), to_apply=callme
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo));
+
+  Cast<HloCallInstruction>(
+      module->entry_computation()->GetInstructionWithName("mycall"))
+      ->set_output_to_operand_aliasing({{{}, {1, {}}}});
+
+  auto status = verifier().Run(module.get()).status();
+  ASSERT_FALSE(status.ok());
+  EXPECT_THAT(status.message(), HasSubstr("Invalid aliasing operand index."));
+}
+
+TEST_F(HloVerifierTest,
+       CheckCallOperandOutputAliasingInvalidOperandShapeIndex) {
+  constexpr absl::string_view hlo = R"(
+    HloModule Module
+
+    callme {
+      ROOT param = (s32[], f32[4]) parameter(0)
+    }
+
+    ENTRY entry {
+      p0 = (s32[], f32[4]) parameter(0)
+      ROOT mycall = (s32[], f32[4]) call(p0), to_apply=callme
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo));
+
+  Cast<HloCallInstruction>(
+      module->entry_computation()->GetInstructionWithName("mycall"))
+      ->set_output_to_operand_aliasing({{{0}, {0, {2}}}});
+
   auto status = verifier().Run(module.get()).status();
   ASSERT_FALSE(status.ok());
   EXPECT_THAT(status.message(),
-              HasSubstr("may not have an output-to-operand aliasing set."));
+              HasSubstr("Invalid aliasing operand shape index."));
+}
+
+TEST_F(HloVerifierTest, CheckCallOperandOutputAliasingInvalidOutputShapeIndex) {
+  constexpr absl::string_view hlo = R"(
+    HloModule Module
+
+    callme {
+      ROOT param = (s32[], f32[4]) parameter(0)
+    }
+
+    ENTRY entry {
+      p0 = (s32[], f32[4]) parameter(0)
+      ROOT mycall = (s32[], f32[4]) call(p0), to_apply=callme
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo));
+
+  Cast<HloCallInstruction>(
+      module->entry_computation()->GetInstructionWithName("mycall"))
+      ->set_output_to_operand_aliasing({{{2}, {0, {0}}}});
+
+  auto status = verifier().Run(module.get()).status();
+  ASSERT_FALSE(status.ok());
+  EXPECT_THAT(status.message(),
+              HasSubstr("Invalid aliasing output shape index."));
+}
+
+TEST_F(HloVerifierTest, CheckCallOperandOutputAliasingShapeMismatch) {
+  constexpr absl::string_view hlo = R"(
+    HloModule Module
+
+    callme {
+      ROOT param = (s32[], f32[4]) parameter(0)
+    }
+
+    ENTRY entry {
+      p0 = (s32[], f32[4]) parameter(0)
+      ROOT mycall = (s32[], f32[4]) call(p0), to_apply=callme
+    }
+  )";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo));
+
+  Cast<HloCallInstruction>(
+      module->entry_computation()->GetInstructionWithName("mycall"))
+      ->set_output_to_operand_aliasing({{{0}, {0, {1}}}});
+
+  auto status = verifier().Run(module.get()).status();
+  ASSERT_FALSE(status.ok());
+  EXPECT_THAT(status.message(),
+              HasSubstr("Different aliasing shapes: f32[4] vs s32[]"));
 }
 
 TEST_F(HloVerifierTest, CheckCustomCallOperandOutputAliasing) {
@@ -1238,6 +1361,33 @@ TEST_F(HloVerifierTestLayoutSensitive,
           "Different aliasing shapes: f32[32,32]{0,1} vs f32[32,32]{1,0}"));
 }
 
+TEST_F(HloVerifierTestLayoutSensitive, VerifyCallAliasConfigLayoutMismatch) {
+  const char* const hlo_string = R"(
+  HloModule module
+
+  callme {
+    p0 = f32[32,32]{0,1} parameter(0)
+    p1 = f32[32,32]{1,0} parameter(1)
+    ROOT tuple = (f32[32,32]{1,0}) tuple(p1)
+  }
+
+  ENTRY main {
+    p0 = f32[32,32]{0,1} parameter(0)
+    p1 = f32[32,32]{1,0} parameter(1)
+    ROOT call = (f32[32,32]{1,0}) call(p0, p1), to_apply=callme
+  })";
+  ASSERT_OK_AND_ASSIGN(auto module, ParseAndReturnUnverifiedModule(hlo_string));
+  Cast<HloCallInstruction>(
+      module->entry_computation()->GetInstructionWithName("call"))
+      ->set_output_to_operand_aliasing({{{0}, {0, {}}}});
+  auto status = verifier().Run(module.get()).status();
+  EXPECT_FALSE(status.ok());
+  EXPECT_THAT(
+      status.message(),
+      HasSubstr(
+          "Different aliasing shapes: f32[32,32]{0,1} vs f32[32,32]{1,0}"));
+}
+
 TEST_F(HloVerifierTestLayoutSensitive,
        VerifyAsyncStartAliasConfigOperandMatchesButComputationMismatches) {
   const char* const hlo_string = R"(
diff --git a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc
index 447f8f3e8895f2..abf40b148bf0f9 100644
--- a/third_party/xla/xla/service/memory_space_assignment/algorithm.cc
+++ b/third_party/xla/xla/service/memory_space_assignment/algorithm.cc
@@ -199,8 +199,8 @@ bool HasAsyncPipelinedWhileLoops(const HloModule& module) {
   return false;
 }
 
-// Returns true if 'buffer' is defined by or used by an async pipelined while
-// loop.
+// Returns true if 'buffer' is defined by, used by, or contained within an
+// async pipelined while loop computation.
 bool IsBufferAliasedToAsyncPipelinedWhileLoop(const HloValue* buffer) {
   if (buffer == nullptr) {
     return false;
@@ -209,11 +209,27 @@ bool IsBufferAliasedToAsyncPipelinedWhileLoop(const HloValue* buffer) {
     if (IsAsyncPipelinedWhileLoop(pos.instruction)) {
       return true;
     }
+    if (pos.instruction->parent() != nullptr) {
+      for (const HloInstruction* caller :
+           pos.instruction->parent()->caller_instructions(HloOpcode::kWhile)) {
+        if (IsAsyncPipelinedWhileLoop(caller)) {
+          return true;
+        }
+      }
+    }
   }
   for (const HloUse& use : buffer->GetUses()) {
     if (IsAsyncPipelinedWhileLoop(use.instruction)) {
       return true;
     }
+    if (use.instruction->parent() != nullptr) {
+      for (const HloInstruction* caller :
+           use.instruction->parent()->caller_instructions(HloOpcode::kWhile)) {
+        if (IsAsyncPipelinedWhileLoop(caller)) {
+          return true;
+        }
+      }
+    }
   }
   return false;
 }
@@ -526,6 +542,41 @@ HloInstruction* GetWhileForBodyRoot(HloInstruction* body_root) {
   return nullptr;
 }
 
+// Returns true if the position corresponds to loop-carried state (parameter or
+// body root tuple) of an async pipelined while loop, or while/DUS/DS in an
+// async pipelined while loop.
+bool IsAsyncPipelinedWhilePosition(const HloPosition& pos) {
+  if (pos.instruction->opcode() == HloOpcode::kParameter) {
+    for (const HloInstruction* caller :
+         pos.instruction->parent()->caller_instructions(HloOpcode::kWhile)) {
+      if (IsAsyncPipelinedWhileLoop(caller)) {
+        return true;
+      }
+    }
+    return false;
+  }
+  if (pos.instruction->opcode() == HloOpcode::kWhile) {
+    return IsAsyncPipelinedWhileLoop(pos.instruction);
+  }
+  if (pos.instruction->opcode() == HloOpcode::kDynamicUpdateSlice ||
+      pos.instruction->opcode() == HloOpcode::kDynamicSlice) {
+    for (const HloInstruction* caller :
+         pos.instruction->parent()->caller_instructions(HloOpcode::kWhile)) {
+      if (IsAsyncPipelinedWhileLoop(caller)) {
+        return true;
+      }
+    }
+    return false;
+  }
+  HloComputation* comp = pos.instruction->parent();
+  if (comp != nullptr && pos.instruction == comp->root_instruction() &&
+      pos.instruction->opcode() == HloOpcode::kTuple) {
+    HloInstruction* while_caller = GetWhileForBodyRoot(pos.instruction);
+    return while_caller != nullptr && IsAsyncPipelinedWhileLoop(while_caller);
+  }
+  return false;
+}
+
 // Returns true if the computation is the body or condition of a while loop.
 bool IsWhileBodyOrConditionComputation(const HloComputation* computation) {
   if (computation == nullptr) {
@@ -672,14 +723,6 @@ MsaAlgorithm::MsaAlgorithm(HloModule* module, AllocationSequence* allocations,
 
 namespace {
 
-bool IsAsyncDynamicSliceOrDynamicUpdateSlice(
-    const HloInstruction* instruction) {
-  return instruction->IsAsynchronous() &&
-         (instruction->async_wrapped_opcode() == HloOpcode::kDynamicSlice ||
-          instruction->async_wrapped_opcode() ==
-              HloOpcode::kDynamicUpdateSlice);
-}
-
 // Finds the allocation value for operand 0 of async updates and dones. We make
 // sure these uses are associated with the AllocationValue whose position
 // is the direct source of the use.
@@ -687,10 +730,10 @@ AllocationValue* FindAllocationValueForAsyncOperationStateUse(
     const HloUse& use,
     absl::Span<AllocationValue> candidate_allocation_values) {
   CHECK_EQ(use.operand_number, 0);
+  HloPosition source_position = GetNonTrivialSourcePosition(
+      HloPosition{use.instruction->mutable_operand(0), use.operand_index});
   for (AllocationValue& allocation_value : candidate_allocation_values) {
-    if (allocation_value.defining_instruction() ==
-            use.instruction->operand(0) &&
-        allocation_value.defining_position().index == use.operand_index) {
+    if (allocation_value.defining_position() == source_position) {
       return &allocation_value;
     }
   }
@@ -755,14 +798,7 @@ AllocationValue* FindLatestAllocationValueForUse(
 AllocationValue* MsaAlgorithm::FindAllocationValueForUse(
     const HloUse& use, absl::Span<AllocationValue> candidate_allocation_values,
     int64_t use_time) const {
-  // Asynchronous DynamicSlice and DynamicUpdateSlice instructions have
-  // specialized liveness and buffer aliasing semantics that differ from
-  // communication async operations (such as AllReduce or AllGather). Do not
-  // restrict their uses to strict start/done pairs.
-  bool is_async_dus_done =
-      IsAsyncDynamicSliceOrDynamicUpdateSlice(use.instruction);
-
-  if (IsAsyncOperationStateUse(use) && !is_async_dus_done) {
+  if (IsAsyncOperationStateUse(use)) {
     // Case A: Async operation state uses. Note, this case covers uses of
     // bundled async state in traditional serial async chains and in pipelined
     // while loops.
@@ -775,19 +811,19 @@ AllocationValue* MsaAlgorithm::FindAllocationValueForUse(
   const absl::flat_hash_map<const HloInstruction*, int64_t>&
       instruction_schedule = hlo_live_range_.instruction_schedule();
 
-  bool is_pipelined_while_async_start_use =
-      has_async_pipelined_while_loops_ &&
-      (use.instruction->opcode() == HloOpcode::kWhile ||
-       (use.instruction->opcode() == HloOpcode::kTuple &&
-        GetWhileForBodyRoot(use.instruction) != nullptr));
+  HloInstruction* while_for_body_root = GetWhileForBodyRoot(use.instruction);
 
-  bool allow_async_state_definitions =
-      is_async_dus_done || is_pipelined_while_async_start_use;
+  bool is_pipelined_while_async_start_use =
+      (use.instruction->opcode() == HloOpcode::kWhile &&
+       IsAsyncPipelinedWhileLoop(use.instruction)) ||
+      (while_for_body_root != nullptr &&
+       IsAsyncPipelinedWhileLoop(while_for_body_root) &&
+       use.instruction->opcode() == HloOpcode::kTuple);
 
   return FindLatestAllocationValueForUse(
       candidate_allocation_values, use_time, use.instruction->parent(),
       instruction_schedule,
-      /*skip_async_state_definitions=*/!allow_async_state_definitions);
+      /*skip_async_state_definitions=*/!is_pipelined_while_async_start_use);
 }
 
 void MsaAlgorithm::CreateAllocationValues(
@@ -877,10 +913,13 @@ void MsaAlgorithm::CreateAllocationValues(
     bool is_async_operation_state =
         IsAsyncOperationStateDefinition(
             allocation_value.defining_instruction()) ||
-        absl::c_any_of(allocation_value.uses(),
-                       [](const AllocationValue::Use& use) {
-                         return IsAsyncOperationStateUse(use.hlo_use);
-                       });
+        (has_async_pipelined_while_loops_ &&
+         IsBufferAliasedToAsyncPipelinedWhileLoop(allocation_value.value()) &&
+         (IsAsyncPipelinedWhilePosition(allocation_value.defining_position()) ||
+          absl::c_any_of(allocation_value.uses(),
+                         [](const AllocationValue::Use& use) {
+                           return IsAsyncOperationStateUse(use.hlo_use);
+                         })));
     // Requiring contiguous allocation ensures that buffers representing inputs
     // and outputs to the async computation maintain temporal contiguity
     // (preventing MSA from evicting them to default memory).
@@ -1031,6 +1070,7 @@ void MsaAlgorithm::FindAliases(
         HloInstruction* while_instruction =
             GetWhileForBodyRoot(use.hlo_use.instruction);
         if (while_instruction != nullptr &&
+            IsAsyncPipelinedWhileLoop(while_instruction) &&
             use.hlo_use.instruction->opcode() == HloOpcode::kTuple &&
             while_instruction->while_body()->num_parameters() > 0) {
           ShapeIndex index = use.hlo_use.operand_index;
@@ -2418,18 +2458,28 @@ MsaAlgorithm::GetContiguousLiveRangesForBuffer(const HloBuffer* buffer) const {
       ShapeIndex operand_index = use.operand_index;
       HloPosition source_position =
           GetNonTrivialSourcePosition(HloPosition{operand, operand_index});
+      bool is_pipelined_async =
+          (has_async_pipelined_while_loops_ &&
+           IsBufferAliasedToAsyncPipelinedWhileLoop(value) &&
+           (IsAsyncOperationStateDefinition(source_position.instruction) ||
+            IsAsyncPipelinedWhilePosition(source_position) ||
+            IsAsyncOperationStateUse(use)));
       if (options_.position_requires_contiguous_allocation_fn(
               source_position) ||
-          IsAsyncOperationStateDefinition(source_position.instruction) ||
-          IsAsyncOperationStateUse(use)) {
+          is_pipelined_async) {
         VLOG(3) << "Adding use " << use.ToString() << " to contiguous position "
                 << source_position.ToString();
         contiguous_positions_to_uses[source_position].push_back(use);
       }
     }
     for (const HloPosition& position : value->positions()) {
+      bool is_pipelined_async_pos =
+          (has_async_pipelined_while_loops_ &&
+           IsBufferAliasedToAsyncPipelinedWhileLoop(value) &&
+           (IsAsyncOperationStateDefinition(position.instruction) ||
+            IsAsyncPipelinedWhilePosition(position)));
       if (options_.position_requires_contiguous_allocation_fn(position) ||
-          IsAsyncOperationStateDefinition(position.instruction)) {
+          is_pipelined_async_pos) {
         if (!(contiguous_positions_to_uses.contains(position))) {
           LOG(WARNING) << "Position " << position.ToString()
                        << " is required to be contiguous but has no uses, "
@@ -2450,9 +2500,13 @@ MsaAlgorithm::GetContiguousLiveRangesForBuffer(const HloBuffer* buffer) const {
         position.instruction->parent() != nullptr &&
         position.instruction->parent()->root_instruction() != nullptr &&
         IsWhileBodyOrConditionComputation(position.instruction->parent())) {
-      end_time = std::max(
-          end_time, hlo_live_range_.instruction_schedule().at(
-                        position.instruction->parent()->root_instruction()));
+      HloInstruction* while_caller = GetWhileForBodyRoot(
+          position.instruction->parent()->root_instruction());
+      if (while_caller != nullptr && IsAsyncPipelinedWhileLoop(while_caller)) {
+        end_time = std::max(
+            end_time, hlo_live_range_.instruction_schedule().at(
+                          position.instruction->parent()->root_instruction()));
+      }
     }
     for (const HloUse& use : uses) {
       end_time = std::max(end_time, GetCorrectedUseTime(use));
@@ -5683,18 +5737,6 @@ absl::StatusOr<AllocationResult> MsaAlgorithm::AllocateAllocationValues(
         // (including loop input operands, loop parameters, and loop body root
         // tuple elements) adopt identical starting offsets in alternate memory.
         AliasedOffset* preferred_offset = nullptr;
-        auto is_while_tuple_position = [](const HloPosition& pos) {
-          if (pos.instruction->opcode() == HloOpcode::kParameter ||
-              pos.instruction->opcode() == HloOpcode::kWhile) {
-            return true;
-          }
-          HloComputation* comp = pos.instruction->parent();
-          if (comp != nullptr && pos.instruction == comp->root_instruction() &&
-              pos.instruction->opcode() == HloOpcode::kTuple) {
-            return true;
-          }
-          return false;
-        };
         if (has_async_pipelined_while_loops_) {
           const HloBuffer& hlo_buffer = alias_analysis_.GetUniqueBufferAt(
               allocation_value_to_update.defining_position().instruction,
@@ -5705,7 +5747,7 @@ absl::StatusOr<AllocationResult> MsaAlgorithm::AllocateAllocationValues(
             preferred_offset = buf_it->second;
           } else {
             for (const HloPosition& position : hlo_buffer.ComputePositions()) {
-              if (!is_while_tuple_position(position)) {
+              if (!IsAsyncPipelinedWhilePosition(position)) {
                 continue;
               }
               HloComputation* comp = position.instruction->parent();
@@ -5723,13 +5765,12 @@ absl::StatusOr<AllocationResult> MsaAlgorithm::AllocateAllocationValues(
           }
         }
         if (preferred_offset == nullptr) {
-          if (!has_async_pipelined_while_loops_) {
-            auto comp_it = preferred_offset_for_computation.find(
-                allocation_value_to_update.computation());
-            if (comp_it != preferred_offset_for_computation.end()) {
-              preferred_offset = comp_it->second;
-            }
-          } else if (is_while_tuple_position(
+          auto comp_it = preferred_offset_for_computation.find(
+              allocation_value_to_update.computation());
+          if (comp_it != preferred_offset_for_computation.end()) {
+            preferred_offset = comp_it->second;
+          } else if (has_async_pipelined_while_loops_ &&
+                     IsAsyncPipelinedWhilePosition(
                          allocation_value_to_update.defining_position())) {
             auto comp_it =
                 pipelined_while_preferred_offset_for_computation_.find(
@@ -6434,7 +6475,7 @@ void MsaAlgorithm::MaybeCreateMirroredParentAllocationForWhileUse(
   // Special case for while loops since the root offset must agree with
   // other offsets: remember the preferred offset for the while loop body.
   AliasedOffset* offset = GetAliasedOffset(*aliased_allocation);
-  if (has_async_pipelined_while_loops_) {
+  if (IsAsyncPipelinedWhileLoop(hlo_use.instruction)) {
     SynchronizeAliasedWhileLoopOffsets(hlo_use, *aliased_allocation, offset);
   } else {
     preferred_offset_for_computation[hlo_use.instruction->while_body()] =
diff --git a/third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc b/third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc
index 69fbaef0a48818..a36ff25dc0d6e9 100644
--- a/third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc
+++ b/third_party/xla/xla/tsl/profiler/utils/xplane_schema.cc
@@ -426,7 +426,9 @@ const StatTypeMap& GetStatTypeMap() {
        {"hbm_power_events", kHbmPowerEvents},
        {"transaction_with_chip_core_id", kTransactionWithChipCoreId},
        {"program_counter", kProgramCounter},
-       {"uses_ici", kUsesIci}});
+       {"uses_ici", kUsesIci},
+       {"dims", kDimensions},
+       {"type", kType}});
   DCHECK_EQ(stat_type_map->size(), kNumStatTypes);
   return *stat_type_map;
 }
diff --git a/third_party/xla/xla/tsl/profiler/utils/xplane_schema.h b/third_party/xla/xla/tsl/profiler/utils/xplane_schema.h
index 6aa8200ed1a8be..3f776c23721ae0 100644
--- a/third_party/xla/xla/tsl/profiler/utils/xplane_schema.h
+++ b/third_party/xla/xla/tsl/profiler/utils/xplane_schema.h
@@ -408,11 +408,13 @@ enum StatType {
   // Program Counter in Oci Descriptors, etc
   kProgramCounter,
   kUsesIci,
+  kDimensions,
+  kType,
   // LINT.ThenChange(:last_stat_type)
 
   // LINT.IfChange(last_stat_type)
   // Change this to point to the last stat type when adding a new one.
-  kLastStatType = kUsesIci,
+  kLastStatType = kType,
   // LINT.ThenChange(:stat_type_enum)
 };