diff --git a/tensorflow/compiler/tests/ternary_ops_test.py b/tensorflow/compiler/tests/ternary_ops_test.py
index 101ca75f8b68be..8e3dad56d9f9c1 100644
--- a/tensorflow/compiler/tests/ternary_ops_test.py
+++ b/tensorflow/compiler/tests/ternary_ops_test.py
@@ -60,18 +60,21 @@ def testLinspace(self, start, end, num):
     self.assertEqual(result[0], expected[0])
 
   def testRange(self):
-    self._testTernary(
-        math_ops.range,
-        np.int32(1),
-        np.int32(2),
-        np.int32(1),
-        expected=np.array([1], dtype=np.int32))
-    self._testTernary(
-        math_ops.range,
-        np.int32(1),
-        np.int32(7),
-        np.int32(2),
-        expected=np.array([1, 3, 5], dtype=np.int32))
+    for dtype in (self.int_types | self.float_types) - {np.uint8}:
+      self._testTernary(
+          math_ops.range,
+          dtype(1),
+          dtype(2),
+          dtype(1),
+          expected=np.array([1], dtype=dtype),
+      )
+      self._testTernary(
+          math_ops.range,
+          dtype(1),
+          dtype(7),
+          dtype(2),
+          expected=np.array([1, 3, 5], dtype=dtype),
+      )
 
   def testSelect(self):
     for dtype in self.numeric_types:
diff --git a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
index d24d1688d188a6..ad2129611fc099 100644
--- a/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
+++ b/tensorflow/compiler/tf2xla/kernels/sequence_ops.cc
@@ -48,10 +48,10 @@ absl::StatusOr<xla::XlaOp> CreateRangeTensor(
   T limit = limit_literal.Get<T>({});
   T delta = delta_literal.Get<T>({});
 
-  if (delta == 0) {
+  if (delta == static_cast<T>(0)) {
     return errors::InvalidArgument("Requires delta != 0: ", delta);
   }
-  if (delta > 0) {
+  if (delta > static_cast<T>(0)) {
     if (start > limit) {
       return errors::InvalidArgument(
           "Requires start <= limit when delta > 0: ", start, "/", limit);
@@ -62,13 +62,21 @@ absl::StatusOr<xla::XlaOp> CreateRangeTensor(
           "Requires start >= limit when delta < 0: ", start, "/", limit);
     }
   }
-  int64_t size =
-      (std::is_integral<T>::value
-           ? static_cast<T>(
-                 limit == start
-                     ? 0
-                     : (std::abs(limit - start) - 1) / std::abs(delta) + 1)
-           : std::ceil(std::abs((limit - start) / delta)));
+  int64_t size;
+  if constexpr (std::is_integral<T>::value) {
+    int64_t start_i = static_cast<int64_t>(start);
+    int64_t limit_i = static_cast<int64_t>(limit);
+    int64_t delta_i = static_cast<int64_t>(delta);
+    size = (limit_i == start_i
+                ? 0
+                : (std::abs(limit_i - start_i) - 1) / std::abs(delta_i) + 1);
+  } else {
+    double start_f = static_cast<double>(start);
+    double limit_f = static_cast<double>(limit);
+    double delta_f = static_cast<double>(delta);
+    size = static_cast<int64_t>(
+        std::ceil(std::abs((limit_f - start_f) / delta_f)));
+  }
 
   return xla::ConstantR0(builder, start) +
          xla::ConstantR0(builder, delta) *
@@ -103,6 +111,13 @@ class RangeOp : public XlaOpKernel {
     DataType type = input_type(0);
     absl::StatusOr<xla::XlaOp> output;
     switch (type) {
+      case DT_INT8:
+        output = CreateRangeTensor<int8_t>(start, limit, delta, ctx->builder());
+        break;
+      case DT_INT16:
+        output =
+            CreateRangeTensor<int16_t>(start, limit, delta, ctx->builder());
+        break;
       case DT_INT32:
         output =
             CreateRangeTensor<int32_t>(start, limit, delta, ctx->builder());
@@ -111,6 +126,26 @@ class RangeOp : public XlaOpKernel {
         output =
             CreateRangeTensor<int64_t>(start, limit, delta, ctx->builder());
         break;
+      case DT_UINT16:
+        output =
+            CreateRangeTensor<uint16_t>(start, limit, delta, ctx->builder());
+        break;
+      case DT_UINT32:
+        output =
+            CreateRangeTensor<uint32_t>(start, limit, delta, ctx->builder());
+        break;
+      case DT_UINT64:
+        output =
+            CreateRangeTensor<uint64_t>(start, limit, delta, ctx->builder());
+        break;
+      case DT_HALF:
+        output =
+            CreateRangeTensor<Eigen::half>(start, limit, delta, ctx->builder());
+        break;
+      case DT_BFLOAT16:
+        output =
+            CreateRangeTensor<bfloat16>(start, limit, delta, ctx->builder());
+        break;
       case DT_FLOAT:
         output = CreateRangeTensor<float>(start, limit, delta, ctx->builder());
         break;
@@ -133,7 +168,7 @@ class RangeOp : public XlaOpKernel {
       xla::XlaOp delta = ctx->Input(2);
       xla::XlaOp limit = ctx->Input(1);
       xla::XlaOp start = ctx->Input(0);
-      if (type == DT_INT32 || type == DT_INT64) {
+      if (DataTypeIsInteger(type)) {
         auto dynamic_size = (xla::Abs(limit - start) + xla::Abs(delta) -
                              xla::One(ctx->builder(), ctx->input_xla_type(0))) /
                             xla::Abs(delta);
diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index 3c5b008c1a74fb..7d1973e434b200 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -95,6 +95,7 @@ tf_cc_test(
     deps = [
         ":cache_dataset_ops",
         ":iterator_ops",
+        ":range_dataset_op",
         ":tensor_slice_dataset_op",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
diff --git a/tensorflow/core/kernels/data/cache_dataset_ops.cc b/tensorflow/core/kernels/data/cache_dataset_ops.cc
index 9bfcbe60a7ffb1..bdda4d13db04f3 100644
--- a/tensorflow/core/kernels/data/cache_dataset_ops.cc
+++ b/tensorflow/core/kernels/data/cache_dataset_ops.cc
@@ -75,9 +75,15 @@ constexpr char kCacheDataset[] = "CacheDataset";
 constexpr char kIncompleteCacheErrorMessage[] =
     "The calling iterator did not fully read the dataset being cached. In "
     "order to avoid unexpected truncation of the dataset, the partially cached "
-    "contents of the dataset  will be discarded. This can happen if you have "
-    "an input pipeline similar to `dataset.cache().take(k).repeat()`. You "
-    "should use `dataset.take(k).cache().repeat()` instead.";
+    "contents of the dataset will be discarded. This can happen if you have "
+    "an input pipeline similar to `dataset.cache().take(k).repeat()`, or if "
+    "downstream operations drop elements (e.g. `batch(drop_remainder=True)`). "
+    "You should use `dataset.take(k).cache().repeat()` instead, or ensure the "
+    "dataset size is a multiple of the batch size before caching. Another "
+    "common workaround is to place the `.cache()` operation after the "
+    "operation that drops elements (like `.batch(...)`), if caching the "
+    "transformed data is acceptable.";
+constexpr size_t kMaxItems = 10000000;  // 10 million
 }  // namespace
 
 class DatasetRandomAccessCache {
@@ -89,15 +95,15 @@ class DatasetRandomAccessCache {
   // out_tensors with the element at that index.
   absl::Status Get(OpKernelContext* ctx, int64_t index,
                    std::vector<Tensor>* out_tensors) {
+    if (index < 0) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Expected index >= 0; Received index: ", index));
+    }
     if (!iter_resource_) {
       TF_ASSIGN_OR_RETURN(iter_resource_,
                           GetIteratorResourceFromDataset(ctx, input_));
       TF_RETURN_IF_ERROR(iter_resource_->SetIteratorFromDataset(ctx, input_));
     }
-    if (index < 0) {
-      return absl::InvalidArgumentError(
-          absl::StrCat("Expected index >= 0; Received index: ", index));
-    }
     if (index >= static_cast<int64_t>(cache_.size())) {
       TF_RETURN_IF_ERROR(ExtendTempCacheToIndex(index, ctx));
     }
@@ -159,12 +165,25 @@ class IteratorRandomAccessCache {
                        element_position));
     }
 
+    if (static_cast<size_t>(element_position) ==
+            std::numeric_limits<size_t>::max() ||
+        static_cast<size_t>(element_position) >= cache_.max_size()) {
+      return absl::InvalidArgumentError(
+          absl::StrCat("Element position too large or invalid."));
+    }
+
     if (element_position < static_cast<int64_t>(cache_.size()) &&
         !cache_[element_position].empty()) {
       *out_tensors = cache_[element_position];
       return absl::OkStatus();
     }
 
+    if (element_position >= kMaxItems) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Requested element_position ", element_position,
+          " exceeds the maximum allowed cache size of ", kMaxItems));
+    }
+
     TF_RETURN_IF_ERROR(input_->Get(ctx, element_position, out_tensors));
     if (element_position >= static_cast<int64_t>(cache_.size())) {
       cache_.resize(element_position + 1);
@@ -721,7 +740,6 @@ class CacheDatasetOp::FileDatasetBase : public DatasetBase {
   Env* const env_;
   const size_t num_tensors_;
   const size_t tensor_index_padding_size_;
-  static constexpr size_t kMaxItems = 10000000;  // 10 million
   const size_t item_index_padding_size_;
 };  // FileDatasetBase
 
diff --git a/tensorflow/core/kernels/data/cache_dataset_ops_test.cc b/tensorflow/core/kernels/data/cache_dataset_ops_test.cc
index ba6dada3d704e6..799734d7a588c8 100644
--- a/tensorflow/core/kernels/data/cache_dataset_ops_test.cc
+++ b/tensorflow/core/kernels/data/cache_dataset_ops_test.cc
@@ -375,14 +375,51 @@ INSTANTIATE_TEST_CASE_P(CacheDatasetOpTest,
                         ParameterizedIteratorSaveAndRestoreTest,
                         ::testing::ValuesIn(IteratorSaveAndRestoreTestCases()));
 
-TEST_F(CacheDatasetOpTest, NegativeIndexTest) {
-  auto params = CacheDatasetParams3();
+TEST_F(CacheDatasetOpTest, NegativeIndexEarlyRejection) {
+  auto range_dataset_params = RangeDatasetParams(0, 20000000, 1);
+  auto params =
+      CacheDatasetParams(range_dataset_params,
+                         /*filename=*/"",
+                         /*output_dtypes=*/{DT_INT64},
+                         /*output_shapes=*/{PartialTensorShape({})}, kNodeName);
   TF_ASSERT_OK(Initialize(params));
   std::vector<Tensor> out_tensors;
   absl::Status status =
-      dataset_->Get(AnyContext(iterator_ctx_.get()), -1, &out_tensors);
-  EXPECT_TRUE(status.code() == absl::StatusCode::kOutOfRange);
-  EXPECT_EQ(status.message(), "Index out of range [0, 3):-1");
+      dataset_->Get(AnyContext(iterator_ctx_.get()), -1LL, &out_tensors);
+  EXPECT_TRUE(status.code() == absl::StatusCode::kInvalidArgument ||
+              status.code() == absl::StatusCode::kOutOfRange);
+}
+
+TEST_F(CacheDatasetOpTest, LargeIndexTest) {
+  auto range_dataset_params = RangeDatasetParams(0, 20000000, 1);
+  auto params =
+      CacheDatasetParams(range_dataset_params,
+                         /*filename=*/"",
+                         /*output_dtypes=*/{DT_INT64},
+                         /*output_shapes=*/{PartialTensorShape({})}, kNodeName);
+  TF_ASSERT_OK(Initialize(params));
+  std::vector<Tensor> out_tensors;
+  int64_t huge_index = std::numeric_limits<int64_t>::max();
+  absl::Status status =
+      dataset_->Get(AnyContext(iterator_ctx_.get()), huge_index, &out_tensors);
+  EXPECT_TRUE(status.code() == absl::StatusCode::kInvalidArgument ||
+              status.code() == absl::StatusCode::kOutOfRange);
+}
+
+TEST_F(CacheDatasetOpTest, BadAllocCrashTest) {
+  auto range_dataset_params = RangeDatasetParams(0, 20000000, 1);
+  auto params =
+      CacheDatasetParams(range_dataset_params,
+                         /*filename=*/"",
+                         /*output_dtypes=*/{DT_INT64},
+                         /*output_shapes=*/{PartialTensorShape({})}, kNodeName);
+  TF_ASSERT_OK(Initialize(params));
+  std::vector<Tensor> out_tensors;
+  absl::Status status =
+      dataset_->Get(AnyContext(iterator_ctx_.get()), 15000000, &out_tensors);
+  EXPECT_EQ(status.code(), absl::StatusCode::kInvalidArgument);
+  EXPECT_TRUE(absl::StrContains(status.message(),
+                                "exceeds the maximum allowed cache size"));
 }
 
 }  // namespace
diff --git a/tensorflow/core/kernels/dilation_ops.cc b/tensorflow/core/kernels/dilation_ops.cc
index 8919bead32fca5..5e6f34708962ad 100644
--- a/tensorflow/core/kernels/dilation_ops.cc
+++ b/tensorflow/core/kernels/dilation_ops.cc
@@ -53,6 +53,9 @@ void ParseAttributes(OpKernelConstruction* context,
   OP_REQUIRES(context, (*strides)[0] == 1 && (*strides)[3] == 1,
               absl::UnimplementedError(
                   "Stride is only supported across spatial dimensions."));
+  OP_REQUIRES(context, (*strides)[1] >= 1 && (*strides)[2] >= 1,
+              absl::InvalidArgumentError(
+                  "Strides in the spatial dimensions must be >= 1."));
 
   OP_REQUIRES_OK(context, context->GetAttr("rates", rates));
   OP_REQUIRES(context, rates->size() == 4,
@@ -61,6 +64,9 @@ void ParseAttributes(OpKernelConstruction* context,
   OP_REQUIRES(context, (*rates)[0] == 1 && (*rates)[3] == 1,
               absl::UnimplementedError(
                   "Rate is only supported across spatial dimensions."));
+  OP_REQUIRES(context, (*rates)[1] >= 1 && (*rates)[2] >= 1,
+              absl::InvalidArgumentError(
+                  "Rates in the spatial dimensions must be >= 1."));
 
   OP_REQUIRES_OK(context, context->GetAttr("padding", padding));
 }
@@ -103,10 +109,12 @@ void ParseSizes(OpKernelContext* context, const std::vector<int32_t>& strides,
 
   // Effective filter size, after introducing rate - 1 zeros between each
   // non-zero filter element.
-  const int filter_rows_eff =
-      filter_rows + (filter_rows - 1) * (*rate_rows - 1);
-  const int filter_cols_eff =
-      filter_cols + (filter_cols - 1) * (*rate_cols - 1);
+  const int64_t filter_rows_eff =
+      static_cast<int64_t>(filter_rows) +
+      static_cast<int64_t>(filter_rows - 1) * (*rate_rows - 1);
+  const int64_t filter_cols_eff =
+      static_cast<int64_t>(filter_cols) +
+      static_cast<int64_t>(filter_cols - 1) * (*rate_cols - 1);
 
   OP_REQUIRES_OK(context, GetWindowedOutputSize(
                               input_rows, filter_rows_eff, /*dilation_rate=*/1,
diff --git a/tensorflow/core/kernels/map_stage_op.cc b/tensorflow/core/kernels/map_stage_op.cc
index e4c674646cb890..cff40deb0b4e9d 100644
--- a/tensorflow/core/kernels/map_stage_op.cc
+++ b/tensorflow/core/kernels/map_stage_op.cc
@@ -448,11 +448,11 @@ class StagingMap : public ResourceBase {
 
     auto it = map_.begin();
 
+    *key = it->first;
+
     TF_RETURN_IF_ERROR(
         copy_or_move_tensors(&it->second, *key, *indices, tuple));
 
-    *key = it->first;
-
     // Remove entry if all the values have been consumed
     if (!std::any_of(
             it->second.begin(), it->second.end(),
diff --git a/tensorflow/core/profiler/utils/BUILD b/tensorflow/core/profiler/utils/BUILD
index 621adf043ee11e..bcc08265beddf0 100644
--- a/tensorflow/core/profiler/utils/BUILD
+++ b/tensorflow/core/profiler/utils/BUILD
@@ -261,7 +261,6 @@ cc_library(
     hdrs = ["hlo_module_utils.h"],
     visibility = internal_visibility([
         "//third_party/odml/model_explorer/backend/adapters/hlo:__pkg__",
-        "//tensorflow/compiler/mlir/lite/experimental/google/tooling/hlo_adapter:__pkg__",
     ]),
     deps = [
         "@org_xprof//xprof/utils:hlo_module_utils",
diff --git a/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py b/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py
index 832fce050ca3d6..70270a53133c2e 100644
--- a/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py
+++ b/tensorflow/python/kernel_tests/data_structures/map_stage_op_test.py
@@ -655,6 +655,67 @@ def testNonScalarKeyMapUnStage(self):
       )
       self.evaluate(v)
 
+  def testMapUnstageNoKeyOutOfRangeIndex(self):
+    # MapUnstageNoKey with an out-of-range index after MapStage must surface
+    # a normal InvalidArgumentError from check_index(), not a fatal CHECK
+    # inside Tensor::CheckIsAlignedAndSingleElement. The CHECK can fire if
+    # popitem() reaches copy_or_move_tensors() with an empty local key
+    # tensor, because check_index() formats its error using
+    # key.scalar<int64_t>()() and Tensor::scalar() requires a 1-element
+    # tensor.
+    stage_op = data_flow_ops.gen_data_flow_ops.map_stage(
+        key=constant_op.constant([1], dtype=dtypes.int64),
+        indices=constant_op.constant([0], dtype=dtypes.int32),
+        values=[constant_op.constant([1.0], dtype=dtypes.float32)],
+        dtypes=[dtypes.float32],
+        capacity=10,
+        memory_limit=0,
+        container='',
+        shared_name='test_map_unstage_no_key_oob',
+        name=None,
+    )
+    self.evaluate(stage_op)
+    with self.assertRaisesRegex(errors.InvalidArgumentError, 'out of bounds'):
+      result = data_flow_ops.gen_data_flow_ops.map_unstage_no_key(
+          indices=[1],
+          dtypes=[dtypes.int64, dtypes.float32],
+          capacity=10,
+          memory_limit=0,
+          container='',
+          shared_name='test_map_unstage_no_key_oob',
+          name=None,
+      )
+      self.evaluate(result)
+
+  def testOrderedMapUnstageNoKeyOutOfRangeIndex(self):
+    # Parallel coverage for the ordered variant. OrderedMapUnstageNoKey
+    # shares StagingMap::popitem() with MapUnstageNoKey via the
+    # StagingMap<bool Ordered> template, so the same out-of-range-index
+    # path must surface InvalidArgumentError rather than abort.
+    stage_op = data_flow_ops.gen_data_flow_ops.ordered_map_stage(
+        key=constant_op.constant([1], dtype=dtypes.int64),
+        indices=constant_op.constant([0], dtype=dtypes.int32),
+        values=[constant_op.constant([1.0], dtype=dtypes.float32)],
+        dtypes=[dtypes.float32],
+        capacity=10,
+        memory_limit=0,
+        container='',
+        shared_name='test_ordered_map_unstage_no_key_oob',
+        name=None,
+    )
+    self.evaluate(stage_op)
+    with self.assertRaisesRegex(errors.InvalidArgumentError, 'out of bounds'):
+      result = data_flow_ops.gen_data_flow_ops.ordered_map_unstage_no_key(
+          indices=[1],
+          dtypes=[dtypes.int64, dtypes.float32],
+          capacity=10,
+          memory_limit=0,
+          container='',
+          shared_name='test_ordered_map_unstage_no_key_oob',
+          name=None,
+      )
+      self.evaluate(result)
+
 
 if __name__ == '__main__':
   test.main()
diff --git a/tensorflow/python/ops/BUILD b/tensorflow/python/ops/BUILD
index fea13cbd711535..16fe6b93b648d0 100644
--- a/tensorflow/python/ops/BUILD
+++ b/tensorflow/python/ops/BUILD
@@ -4832,3 +4832,32 @@ py_test(
         "//tensorflow/python/platform:client_testlib",
     ],
 )
+
+py_library(
+    name = "quantized_dense",
+    srcs = ["quantized_dense.py"],
+    srcs_version = "PY3",
+    deps = [
+        ":array_ops",
+        ":math_ops",
+        ":nn_ops",
+        ":random_ops",
+        ":variables",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/module",
+    ],
+)
+
+tf_py_strict_test(
+    name = "quantized_dense_test",
+    size = "small",
+    srcs = ["quantized_dense_test.py"],
+    python_version = "PY3",
+    deps = [
+        ":quantized_dense",
+        ":random_ops",
+        "//tensorflow/python/framework:random_seed",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
diff --git a/tensorflow/python/ops/math_ops.py b/tensorflow/python/ops/math_ops.py
index c9676bf6fabcf0..eb4dc0a28abe64 100644
--- a/tensorflow/python/ops/math_ops.py
+++ b/tensorflow/python/ops/math_ops.py
@@ -2097,8 +2097,13 @@ def range(start, limit=None, delta=1, dtype=None, name="range"):  # pylint: disa
     # infer dtype if not explicitly provided
     if dtype is None:
       dtype_hierarchy = [
+          dtypes.int8,
+          dtypes.int16,
           dtypes.int32,
           dtypes.int64,
+          dtypes.uint16,
+          dtypes.uint32,
+          dtypes.uint64,
           dtypes.float16,
           dtypes.bfloat16,
           dtypes.float32,
diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops.py b/tensorflow/python/ops/numpy_ops/np_math_ops.py
index 18e885f5dbdc0d..d7ed4722770933 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops.py
@@ -1424,6 +1424,14 @@ def concatenate(arys, axis=0):  # pylint: disable=missing-function-docstring
     )
   dtype = np_utils.result_type(*arys)
   arys = [np_array_ops.array(array, dtype=dtype) for array in arys]
+  if axis is None:
+    # NumPy flattens every input before concatenating when axis is None.
+    # Reshaping an already flat array is a no-op, so skip the op dispatch.
+    arys = [
+        array if array.shape.ndims == 1 else array_ops.reshape(array, [-1])
+        for array in arys
+    ]
+    axis = 0
   return array_ops.concat(arys, axis)
 
 
diff --git a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
index cf7e185cfeaf4e..26f33fd00c3b19 100644
--- a/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
+++ b/tensorflow/python/ops/numpy_ops/np_math_ops_test.py
@@ -667,6 +667,16 @@ def testSignBit(self):
     negative_zero = ops.convert_to_tensor([-0.0], dtype=dtypes.bfloat16)
     self.assertAllEqual(np_math_ops.signbit(negative_zero), [True])
 
+  def testConcatenateAxisNone(self):
+    a = np_array_ops.array([1, 2])
+    b = np_array_ops.array([[3], [4]])
+    self.assertAllEqual(
+        np_math_ops.concatenate([a, b], axis=None), [1, 2, 3, 4]
+    )
+    self.assertAllEqual(
+        np_math_ops.concatenate(np_array_ops.array([[5, 6]]), axis=None), [5, 6]
+    )
+
   def testIsInfFamilyNonFloatInputs(self):
     # A non-floating input has no infinities, but the result must still be an
     # elementwise boolean array shaped like the input, as numpy returns, and
diff --git a/tensorflow/python/ops/quantized_dense.py b/tensorflow/python/ops/quantized_dense.py
new file mode 100644
index 00000000000000..890e94fe65d872
--- /dev/null
+++ b/tensorflow/python/ops/quantized_dense.py
@@ -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.
+# ==============================================================================
+"""Quantized Dense module."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.module import module
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops import variables
+
+
+class QuantizedDense(module.Module):
+  """A densely-connected layer with weight quantization.
+
+  This module acts like a standard Dense layer but simulates
+  4-bit or 8-bit weight quantization using fake quantization nodes.
+  """
+
+  def __init__(self, units, bits=8, use_bias=True, name=None):
+    super(QuantizedDense, self).__init__(name=name)
+    self.units = int(units)
+    self.bits = int(bits)
+    if self.bits not in [4, 8]:
+      raise ValueError("Only 4-bit and 8-bit quantization are supported.")
+    self.use_bias = use_bias
+    self.kernel = None
+    self.bias = None
+
+  def __call__(self, inputs):
+    inputs = ops.convert_to_tensor(inputs)
+    if self.kernel is None:
+      last_dim = inputs.shape[-1]
+      if last_dim is None:
+        raise ValueError(
+            "The last dimension of the inputs to `QuantizedDense` should be"
+            " defined."
+        )
+      last_dim = int(last_dim)
+      # Initialize weights with glorot uniform
+      limit = math_ops.sqrt(6.0 / (last_dim + self.units))
+      self.kernel = variables.Variable(
+          initial_value=random_ops.random_uniform(
+              [last_dim, self.units],
+              minval=-limit,
+              maxval=limit,
+              dtype=inputs.dtype,
+          ),
+          name="kernel",
+          trainable=True,
+      )
+      if self.use_bias:
+        self.bias = variables.Variable(
+            initial_value=array_ops.zeros(
+                [
+                    self.units,
+                ],
+                dtype=inputs.dtype,
+            ),
+            name="bias",
+            trainable=True,
+        )
+
+    kernel = math_ops.cast(self.kernel, dtypes.float32)
+    min_val = math_ops.reduce_min(kernel)
+    max_val = math_ops.reduce_max(kernel)
+    max_val = math_ops.maximum(max_val, min_val + 1e-5)
+
+    quantized_kernel = array_ops.fake_quant_with_min_max_vars(
+        kernel, min_val, max_val, num_bits=self.bits, narrow_range=True
+    )
+    quantized_kernel = math_ops.cast(quantized_kernel, inputs.dtype)
+
+    rank = inputs.shape.rank
+    if rank is not None and rank <= 2:
+      outputs = math_ops.matmul(a=inputs, b=quantized_kernel)
+    else:
+      outputs = math_ops.tensordot(
+          inputs, quantized_kernel, [[rank - 1 if rank else -1], [0]]
+      )
+
+    if self.use_bias:
+      outputs = nn_ops.bias_add(outputs, self.bias)
+
+    return outputs
diff --git a/tensorflow/python/ops/quantized_dense_test.py b/tensorflow/python/ops/quantized_dense_test.py
new file mode 100644
index 00000000000000..8734cc3aae55cc
--- /dev/null
+++ b/tensorflow/python/ops/quantized_dense_test.py
@@ -0,0 +1,58 @@
+# 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.
+# ==============================================================================
+"""Tests for quantized Dense module."""
+
+from tensorflow.python.framework import random_seed
+from tensorflow.python.ops import random_ops
+from tensorflow.python.ops.quantized_dense import QuantizedDense
+from tensorflow.python.platform import test
+
+
+class QuantizedDenseTest(test.TestCase):
+
+  def setUp(self):
+    super(QuantizedDenseTest, self).setUp()
+    random_seed.set_random_seed(0)
+
+  def test_quantized_dense_basic(self):
+    inputs = random_ops.random_uniform((32, 128))
+
+    # Test 8-bit quantization
+    layer_8bit = QuantizedDense(64, bits=8)
+    out_8bit = layer_8bit(inputs)
+
+    self.assertEqual(out_8bit.shape, (32, 64))
+    self.assertEqual(layer_8bit.kernel.shape, (128, 64))
+    self.assertEqual(layer_8bit.bias.shape, (64,))
+
+  def test_quantized_dense_4bit(self):
+    inputs = random_ops.random_uniform((16, 32))
+
+    # Test 4-bit quantization
+    layer_4bit = QuantizedDense(16, bits=4, use_bias=False)
+    out_4bit = layer_4bit(inputs)
+
+    self.assertEqual(out_4bit.shape, (16, 16))
+    self.assertIsNone(layer_4bit.bias)
+
+  def test_invalid_bits(self):
+    with self.assertRaisesRegex(
+        ValueError, "Only 4-bit and 8-bit quantization"
+    ):
+      QuantizedDense(32, bits=16)
+
+
+if __name__ == "__main__":
+  test.main()
diff --git a/third_party/flatbuffers/workspace.bzl b/third_party/flatbuffers/workspace.bzl
index aa0ba40e8aec9d..89aa1bdc59ce74 100644
--- a/third_party/flatbuffers/workspace.bzl
+++ b/third_party/flatbuffers/workspace.bzl
@@ -14,7 +14,7 @@ def repo():
         sha256 = _FLATBUFFERS_SHA256,
         urls = tf_mirror_urls("https://github.com/google/flatbuffers/archive/v%s.tar.gz" % _FLATBUFFERS_VERSION),
         build_file = "//third_party/flatbuffers:flatbuffers.BUILD",
-        system_build_file = "//third_party/flatbuffers:BUILD.system",
+        system_build_file = "//third_party/systemlibs:flatbuffers.BUILD",
         link_files = {
             "//third_party/flatbuffers:build_defs.bzl": "build_defs.bzl",
         },
diff --git a/third_party/icu/workspace.bzl b/third_party/icu/workspace.bzl
index 3773cc43ddaeb6..6fb281e548c391 100644
--- a/third_party/icu/workspace.bzl
+++ b/third_party/icu/workspace.bzl
@@ -11,6 +11,7 @@ def repo():
         sha256 = "588e431f77327c39031ffbb8843c0e3bc122c211374485fa87dc5f3faff24061",
         urls = tf_mirror_urls("https://github.com/unicode-org/icu/releases/download/release-77-1/icu4c-77_1-src.tgz"),
         build_file = "//third_party/icu:icu.BUILD",
+        system_build_file = "//third_party/systemlibs:icu.BUILD",
         patch_file = ["//third_party/icu:udata.patch"],
         patch_cmds = [
             "rm -f source/common/BUILD.bazel",
diff --git a/third_party/jpeg/workspace.bzl b/third_party/jpeg/workspace.bzl
index 631cc933bc60d9..579d95ba4fed32 100644
--- a/third_party/jpeg/workspace.bzl
+++ b/third_party/jpeg/workspace.bzl
@@ -9,5 +9,5 @@ def repo():
         sha256 = "a78b05c0d8427a90eb5b4eb08af25309770c8379592bb0b8a863373128e6143f",
         strip_prefix = "libjpeg-turbo-2.1.4",
         build_file = "//third_party/jpeg:jpeg.BUILD",
-        system_build_file = "//third_party/jpeg:BUILD.system",
+        system_build_file = "//third_party/systemlibs:jpeg.BUILD",
     )
diff --git a/third_party/flatbuffers/BUILD.system b/third_party/systemlibs/flatbuffers.BUILD
similarity index 96%
rename from third_party/flatbuffers/BUILD.system
rename to third_party/systemlibs/flatbuffers.BUILD
index 8fe4d7a590719f..b1d63b4ca0fd77 100644
--- a/third_party/flatbuffers/BUILD.system
+++ b/third_party/systemlibs/flatbuffers.BUILD
@@ -1,7 +1,7 @@
 licenses(["notice"])  # Apache 2.0
 
 filegroup(
-    name = "LICENSE.txt",
+    name = "LICENSE",
     visibility = ["//visibility:public"],
 )
 
diff --git a/third_party/icu/BUILD.system b/third_party/systemlibs/icu.BUILD
similarity index 100%
rename from third_party/icu/BUILD.system
rename to third_party/systemlibs/icu.BUILD
diff --git a/third_party/jpeg/BUILD.system b/third_party/systemlibs/jpeg.BUILD
similarity index 100%
rename from third_party/jpeg/BUILD.system
rename to third_party/systemlibs/jpeg.BUILD
diff --git a/third_party/systemlibs/pybind11.BUILD b/third_party/systemlibs/pybind11.BUILD
index 711fcf13b5d863..e90af0bb5879d2 100644
--- a/third_party/systemlibs/pybind11.BUILD
+++ b/third_party/systemlibs/pybind11.BUILD
@@ -18,6 +18,13 @@ package(default_visibility = ["//visibility:public"])
 cc_library(
     name = "pybind11",
     deps = [
-        "@xla@xla//third_party/python_runtime:headers",
+        "@xla//third_party/python_runtime:headers",
     ],
 )
+
+# Needed by pybind11_bazel.
+config_setting(
+    name = "msvc_compiler",
+    flag_values = {"@bazel_tools//tools/cpp:compiler": "msvc-cl"},
+    visibility = ["//visibility:public"],
+)
diff --git a/third_party/xla/MODULE.bazel b/third_party/xla/MODULE.bazel
index 64b3d29e143192..ae3b0122f1d23f 100644
--- a/third_party/xla/MODULE.bazel
+++ b/third_party/xla/MODULE.bazel
@@ -23,7 +23,7 @@ bazel_dep(name = "pybind11_abseil", version = "202402.0")
 bazel_dep(name = "pybind11_bazel", version = "3.0.0")
 bazel_dep(name = "pybind11_protobuf", version = "0.0.0-20250210-f02a2b7")
 bazel_dep(name = "re2", version = "2025-11-05.bcr.1", repo_name = "com_googlesource_code_re2")
-bazel_dep(name = "rules_cc", version = "0.2.18")
+bazel_dep(name = "rules_cc", version = "0.2.20")
 bazel_dep(name = "rules_java", version = "8.16.1")
 bazel_dep(name = "rules_license", version = "1.0.0")
 bazel_dep(name = "rules_python", version = "2.2.0")
@@ -281,7 +281,7 @@ nvshmem_redist = use_extension("@rules_ml_toolchain//extensions:nvshmem_redist.b
 use_repo(nvshmem_redist, "nvidia_nvshmem")
 
 toolchain_ext = use_extension("@rules_ml_toolchain//extensions:toolchain.bzl", "toolchain_ext")
-use_repo(toolchain_ext, "llvm18_linux_x86_64", "llvm_linux_x86_64")
+use_repo(toolchain_ext, "llvm18_linux_x86_64", "llvm_linux_aarch64", "llvm_linux_x86_64")
 
 register_toolchains("@rules_ml_toolchain//cc:linux_x86_64_linux_x86_64")
 
diff --git a/third_party/xla/WORKSPACE b/third_party/xla/WORKSPACE
index 37e4edee2a579f..14cf2eb79476c2 100644
--- a/third_party/xla/WORKSPACE
+++ b/third_party/xla/WORKSPACE
@@ -35,6 +35,10 @@ load("@bazel_features//:deps.bzl", "bazel_features_deps")
 
 bazel_features_deps()
 
+load("@rules_cc//cc:extensions.bzl", "compatibility_proxy_repo")
+
+compatibility_proxy_repo()
+
 # Initialize hermetic C++
 load("@rules_ml_toolchain//cc/deps:cc_toolchain_deps.bzl", "cc_toolchain_deps")
 
diff --git a/third_party/xla/third_party/llvm/build.patch b/third_party/xla/third_party/llvm/build.patch
index 641bdb95bc262c..032406b4e16483 100644
--- a/third_party/xla/third_party/llvm/build.patch
+++ b/third_party/xla/third_party/llvm/build.patch
@@ -15,74 +15,6 @@
 diff --git a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel
 --- a/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel
 +++ b/utils/bazel/llvm-project-overlay/llvm/BUILD.bazel
-@@ -317,19 +317,19 @@
- config_setting(
-     name = "is_windows_clang_mingw",
-     constraint_values = ["@platforms//os:windows"],
--    flag_values = {"@rules_cc//cc/compiler:compiler": "clang"},
-+    flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang"},
- )
- 
- config_setting(
-     name = "is_windows_clang_cl",
-     constraint_values = ["@platforms//os:windows"],
--    flag_values = {"@rules_cc//cc/compiler:compiler": "clang-cl"},
-+    flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang-cl"},
- )
- 
- config_setting(
-     name = "is_windows_msvc",
-     constraint_values = ["@platforms//os:windows"],
--    flag_values = {"@rules_cc//cc/compiler:compiler": "msvc-cl"},
-+    flag_values = {"@rules_cc//cc/private/toolchain:compiler": "msvc-cl"},
- )
- 
- config_setting(
-@@ -338,7 +338,7 @@
-         "@platforms//cpu:aarch64",
-         "@platforms//os:windows",
-     ],
--    flag_values = {"@rules_cc//cc/compiler:compiler": "clang"},
-+    flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang"},
- )
- 
- config_setting(
-@@ -347,7 +347,7 @@
-         "@platforms//cpu:aarch64",
-         "@platforms//os:windows",
-     ],
--    flag_values = {"@rules_cc//cc/compiler:compiler": "clang-cl"},
-+    flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang-cl"},
- )
- 
- config_setting(
-@@ -356,7 +356,7 @@
-         "@platforms//cpu:aarch64",
-         "@platforms//os:windows",
-     ],
--    flag_values = {"@rules_cc//cc/compiler:compiler": "msvc-cl"},
-+    flag_values = {"@rules_cc//cc/private/toolchain:compiler": "msvc-cl"},
- )
- 
- config_setting(
-@@ -365,7 +365,7 @@
-         "@platforms//cpu:x86_64",
-         "@platforms//os:windows",
-     ],
--    flag_values = {"@rules_cc//cc/compiler:compiler": "clang"},
-+    flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang"},
- )
- 
- config_setting(
-@@ -374,7 +374,7 @@
-         "@platforms//cpu:x86_64",
-         "@platforms//os:windows",
-     ],
--    flag_values = {"@rules_cc//cc/compiler:compiler": "clang-cl"},
-+    flag_values = {"@rules_cc//cc/private/toolchain:compiler": "clang-cl"},
- )
- 
- BLAKE3_x86_64_ASM_SOURCE_PATTERNS = [
 @@ -430,7 +430,8 @@
          "@platforms//cpu:aarch64": [
              "lib/Support/BLAKE3/blake3_neon.c",
diff --git a/third_party/xla/third_party/stablehlo/temporary.patch b/third_party/xla/third_party/stablehlo/temporary.patch
index e1b5fa813d250c..ebb9dda1528d7c 100644
--- a/third_party/xla/third_party/stablehlo/temporary.patch
+++ b/third_party/xla/third_party/stablehlo/temporary.patch
@@ -53,295 +53,6 @@ diff --ruN a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
  * `is_type_name(x: Value | Placeholder | Type) -> Value`. Available for all
  types. For example, `is_float(x)` returns `true` if `x` is a `FloatType`.
  If `x` is a value or placeholder, this function is a shortcut for
-diff --ruN a/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp b/stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp
---- stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp
-+++ stablehlo/stablehlo/dialect/ReplicaGroupUtils.cpp
-@@ -20,6 +20,7 @@
- #include <utility>
- 
- #include "llvm/ADT/DenseSet.h"
-+#include "llvm/ADT/STLExtras.h"
- #include "llvm/ADT/SmallVector.h"
- #include "llvm/ADT/StringRef.h"
- #include "mlir/IR/Attributes.h"
-@@ -31,57 +32,53 @@
- namespace mlir {
- namespace stablehlo {
- 
--static SmallVector<SmallVector<int64_t>>
-+namespace {
-+
-+struct ReindexedAxes {
-+  SmallVector<int64_t> splitAxisSizes;
-+  SmallVector<int64_t> groupedAxisIndices;
-+};
-+
-+// Generates replica groups from the reshaped mesh axis sizes and the indices of
-+// the communication axes using a Reshape-Transpose permutation.
-+SmallVector<SmallVector<int64_t>>
- flattenedReplicaGroupsFromTransposePermutation(
--    const SmallVector<StringRef>& meshAxisNames,
--    const SmallVector<StringRef>& commAxisNames,
--    const llvm::DenseSet<StringRef>& commAxisSet,
--    const SmallVector<int64_t>& axisSizes,
--    const SmallVector<int64_t>& deviceIds, int64_t totalDevices) {
--  // Reshape and Transpose equivalence bridging XLA TileAssignment behavior.
-+    ArrayRef<int64_t> axisSizes, ArrayRef<int64_t> groupedAxisIndices,
-+    ArrayRef<int64_t> deviceIds, int64_t totalDevices) {
-+  llvm::DenseSet<int64_t> groupedAxisSet(groupedAxisIndices.begin(),
-+                                         groupedAxisIndices.end());
-   SmallVector<int64_t> transposeAxes;
-   // Non-grouped axes first
--  for (size_t i = 0; i < meshAxisNames.size(); ++i) {
--    if (!commAxisSet.count(meshAxisNames[i])) {
-+  for (size_t i = 0; i < axisSizes.size(); ++i) {
-+    if (!groupedAxisSet.count(i)) {
-       transposeAxes.push_back(i);
-     }
-   }
--  // Grouped axes
--  for (const auto& name : commAxisNames) {
--    for (size_t i = 0; i < meshAxisNames.size(); ++i) {
--      if (meshAxisNames[i] == name) {
--        transposeAxes.push_back(i);
--        break;
--      }
--    }
--  }
--
--  SmallVector<int64_t> transposedSizes(meshAxisNames.size());
--  for (size_t i = 0; i < meshAxisNames.size(); ++i) {
-+  // Grouped axes in the specified order
-+  for (int64_t idx : groupedAxisIndices) {
-+    transposeAxes.push_back(idx);
-+  }
-+
-+  SmallVector<int64_t> transposedSizes(axisSizes.size());
-+  for (size_t i = 0; i < axisSizes.size(); ++i) {
-     transposedSizes[i] = axisSizes[transposeAxes[i]];
-   }
- 
--  // Compute strides for original shape
--  SmallVector<int64_t> originalStrides(meshAxisNames.size(), 1);
--  for (int i = static_cast<int>(meshAxisNames.size()) - 2; i >= 0; --i) {
-+  // Compute strides for reshaped shape
-+  SmallVector<int64_t> originalStrides(axisSizes.size(), 1);
-+  for (int i = static_cast<int>(axisSizes.size()) - 2; i >= 0; --i) {
-     originalStrides[i] = originalStrides[i + 1] * axisSizes[i + 1];
-   }
- 
-   // Compute strides for transposed shape
--  SmallVector<int64_t> transposedStrides(meshAxisNames.size(), 1);
--  for (int i = static_cast<int>(meshAxisNames.size()) - 2; i >= 0; --i) {
-+  SmallVector<int64_t> transposedStrides(axisSizes.size(), 1);
-+  for (int i = static_cast<int>(axisSizes.size()) - 2; i >= 0; --i) {
-     transposedStrides[i] = transposedStrides[i + 1] * transposedSizes[i + 1];
-   }
- 
--  // Generate chunks
-   int64_t numDevicesPerGroup = 1;
--  for (auto name : commAxisNames) {
--    for (size_t i = 0; i < meshAxisNames.size(); ++i) {
--      if (meshAxisNames[i] == name) {
--        numDevicesPerGroup *= axisSizes[i];
--        break;
--      }
--    }
-+  for (int64_t idx : groupedAxisIndices) {
-+    numDevicesPerGroup *= axisSizes[idx];
-   }
-   int64_t numGroups = totalDevices / numDevicesPerGroup;
- 
-@@ -93,7 +90,7 @@
-     for (int64_t j = 0; j < numDevicesPerGroup; ++j) {
-       int64_t linearTransposeIdx = i * numDevicesPerGroup + j;
-       int64_t originalIndex = 0;
--      for (size_t k = 0; k < meshAxisNames.size(); ++k) {
-+      for (size_t k = 0; k < axisSizes.size(); ++k) {
-         int64_t coord =
-             (linearTransposeIdx / transposedStrides[k]) % transposedSizes[k];
-         originalIndex += coord * originalStrides[transposeAxes[k]];
-@@ -102,9 +99,128 @@
-     }
-     groups.push_back(std::move(group));
-   }
--
-   return groups;
- }
-+
-+// Splits mesh axes based on sub-axis references and computes the corresponding
-+// indices for the communication axes.
-+FailureOr<ReindexedAxes> computeReindexedAxes(ArrayRef<MeshAxisAttr> axesInMesh,
-+                                              ArrayAttr commAxes,
-+                                              Location loc) {
-+  ReindexedAxes result;
-+
-+  // Validate commAxes and verify that all mesh axes exist and have valid sizes.
-+  for (auto attr : commAxes) {
-+    auto shloAxisRef = llvm::dyn_cast<AxisRefAttr>(attr);
-+    if (!shloAxisRef) {
-+      return emitError(loc) << "expected AxisRefAttr in comm_axes";
-+    }
-+    StringRef axisName = shloAxisRef.getName();
-+    bool found = false;
-+    for (auto meshAxis : axesInMesh) {
-+      if (meshAxis.getName() == axisName) {
-+        found = true;
-+        if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) {
-+          int64_t preSize = subAxisInfo.getPreSize();
-+          int64_t size = subAxisInfo.getSize();
-+          if (preSize < 1 || size < 1) {
-+            return emitError(loc)
-+                   << "sub-axis pre_size and size must be at least 1";
-+          }
-+          int64_t nextPreSize = preSize * size;
-+          if (nextPreSize > meshAxis.getSize() ||
-+              meshAxis.getSize() % nextPreSize != 0) {
-+            return emitError(loc)
-+                   << "sub-axis (pre_size * size) must divide mesh axis size";
-+          }
-+        }
-+        break;
-+      }
-+    }
-+    if (!found) {
-+      return emitError(loc)
-+             << "axis '" << axisName << "' not found in mesh definition";
-+    }
-+  }
-+
-+  // Split each mesh axis according to the referenced subaxes.
-+  struct SplitDim {
-+    StringRef axisName;
-+    int64_t preSize;
-+    int64_t size;
-+    int64_t dimIndex;
-+  };
-+  SmallVector<SplitDim> splitDims;
-+
-+  for (auto meshAxis : axesInMesh) {
-+    StringRef axisName = meshAxis.getName();
-+    int64_t axisSize = meshAxis.getSize();
-+
-+    SmallVector<int64_t> preSizes = {1, axisSize};
-+    for (auto attr : commAxes) {
-+      auto shloAxisRef = llvm::cast<AxisRefAttr>(attr);
-+      if (shloAxisRef.getName() == axisName) {
-+        if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) {
-+          preSizes.push_back(subAxisInfo.getPreSize());
-+          preSizes.push_back(subAxisInfo.getPreSize() * subAxisInfo.getSize());
-+        }
-+      }
-+    }
-+
-+    llvm::sort(preSizes);
-+    preSizes.erase(llvm::unique(preSizes), preSizes.end());
-+
-+    for (size_t j = 0; j < preSizes.size() - 1; ++j) {
-+      int64_t segPreSize = preSizes[j];
-+      int64_t segSize = preSizes[j + 1] / segPreSize;
-+      int64_t dimIdx = result.splitAxisSizes.size();
-+      result.splitAxisSizes.push_back(segSize);
-+      splitDims.push_back({axisName, segPreSize, segSize, dimIdx});
-+    }
-+  }
-+
-+  // Map each communication axis to its corresponding split dimension.
-+  llvm::DenseSet<int64_t> groupedSet;
-+  for (auto attr : commAxes) {
-+    auto shloAxisRef = llvm::cast<AxisRefAttr>(attr);
-+    StringRef axisName = shloAxisRef.getName();
-+    int64_t reqPreSize = 1;
-+    int64_t reqSize = 0;
-+    if (auto subAxisInfo = shloAxisRef.getSubAxisInfo()) {
-+      reqPreSize = subAxisInfo.getPreSize();
-+      reqSize = subAxisInfo.getSize();
-+    } else {
-+      for (auto meshAxis : axesInMesh) {
-+        if (meshAxis.getName() == axisName) {
-+          reqSize = meshAxis.getSize();
-+          break;
-+        }
-+      }
-+    }
-+
-+    bool matched = false;
-+    for (const auto& splitDim : splitDims) {
-+      if (splitDim.axisName == axisName && splitDim.preSize == reqPreSize &&
-+          splitDim.size == reqSize) {
-+        if (!groupedSet.insert(splitDim.dimIndex).second) {
-+          return emitError(loc)
-+                 << "Duplicate or overlapping communication axis: " << axisName;
-+        }
-+        result.groupedAxisIndices.push_back(splitDim.dimIndex);
-+        matched = true;
-+        break;
-+      }
-+    }
-+    if (!matched) {
-+      return emitError(loc) << "Invalid or overlapping communication axis on '"
-+                            << axisName << "'";
-+    }
-+  }
-+
-+  return result;
-+}
-+
-+}  // namespace
- 
- FailureOr<SmallVector<SmallVector<int64_t>>> flattenReplicaGroupMeshAxes(
-     Attribute meshAttr, ArrayAttr commAxes, Location loc) {
-@@ -120,34 +236,13 @@
-   if (!mesh)
-     return emitOptionalError(loc, "expected stablehlo.mesh for mesh attribute");
- 
--  auto axesInMesh = mesh.getAxes();
--
--  // Identify which axes are communication axes.
--  llvm::SmallVector<StringRef> commAxisNames;
--  llvm::DenseSet<StringRef> commAxisSet;
--  for (auto attr : commAxes) {
--    auto shloAxisRef = llvm::dyn_cast<AxisRefAttr>(attr);
--    if (!shloAxisRef) {
--      return emitError(loc) << "expected AxisRefAttr in comm_axes";
--    }
--    if (shloAxisRef.getSubAxisInfo()) {
--      return emitError(loc) << "Subaxes are not supported in "
--                               "flattenReplicaGroupMeshAxes";
--    }
--    commAxisNames.push_back(shloAxisRef.getName());
--    commAxisSet.insert(shloAxisRef.getName());
--  }
--
--  // Calculate total devices and axis sizes
-+  FailureOr<ReindexedAxes> reindexedAxes =
-+      computeReindexedAxes(mesh.getAxes(), commAxes, loc);
-+  if (failed(reindexedAxes)) return failure();
- 
-   int64_t totalDevices = 1;
--  SmallVector<int64_t> axisSizes;
--  SmallVector<StringRef> meshAxisNames;
--  for (auto meshAxis : axesInMesh) {
--    auto typedMeshAxis = llvm::cast<stablehlo::MeshAxisAttr>(meshAxis);
--    axisSizes.push_back(typedMeshAxis.getSize());
--    meshAxisNames.push_back(typedMeshAxis.getName());
--    totalDevices *= typedMeshAxis.getSize();
-+  for (auto meshAxis : mesh.getAxes()) {
-+    totalDevices *= llvm::cast<stablehlo::MeshAxisAttr>(meshAxis).getSize();
-   }
- 
-   SmallVector<int64_t> deviceIds;
-@@ -160,8 +255,8 @@
-   }
- 
-   return flattenedReplicaGroupsFromTransposePermutation(
--      meshAxisNames, commAxisNames, commAxisSet, axisSizes, deviceIds,
--      totalDevices);
-+      reindexedAxes->splitAxisSizes, reindexedAxes->groupedAxisIndices,
-+      deviceIds, totalDevices);
- }
- 
- }  // namespace stablehlo
 diff --ruN a/stablehlo/stablehlo/dialect/Serialization.cpp b/stablehlo/stablehlo/dialect/Serialization.cpp
 --- stablehlo/stablehlo/dialect/Serialization.cpp
 +++ stablehlo/stablehlo/dialect/Serialization.cpp
@@ -776,196 +487,6 @@ diff --ruN a/stablehlo/stablehlo/tests/chlo/chlo_legalize_to_stablehlo.mlir b/st
  // CHECK-LABEL:   func.func @ragged_dot_mode_3(
  // CHECK-SAME:      %[[ARG0:.*]]: tensor<2x3x5xf32>,
  // CHECK-SAME:      %[[ARG1:.*]]: tensor<2x5x7xf32>,
-diff --ruN a/stablehlo/stablehlo/tests/interpret/all_gather.mlir b/stablehlo/stablehlo/tests/interpret/all_gather.mlir
---- stablehlo/stablehlo/tests/interpret/all_gather.mlir
-+++ stablehlo/stablehlo/tests/interpret/all_gather.mlir
-@@ -133,3 +133,41 @@
-     func.return
-   }
- }
-+
-+// -----
-+
-+module @mesh_axes_subaxis {
-+  func.func @all_gather(%operand : tensor<1xi64>) -> tensor<2xi64> {
-+    %result = "stablehlo.all_gather"(%operand) {
-+      all_gather_dim = 0 : i64,
-+      replica_groups = #stablehlo.replica_group_mesh_axes<mesh = #stablehlo.mesh<axes = [#stablehlo.mesh_axis<name = "x", size = 2>, #stablehlo.mesh_axis<name = "y", size = 4>]>, axes = [#stablehlo.axis_ref<name = "y", sub_axis_info = (1)2>]>
-+    } : (tensor<1xi64>) -> tensor<2xi64>
-+    return %result : tensor<2xi64>
-+  }
-+  func.func @main() {
-+    %p0 = stablehlo.constant dense<[0]> : tensor<1xi64>
-+    %p1 = stablehlo.constant dense<[1]> : tensor<1xi64>
-+    %p2 = stablehlo.constant dense<[2]> : tensor<1xi64>
-+    %p3 = stablehlo.constant dense<[3]> : tensor<1xi64>
-+    %p4 = stablehlo.constant dense<[4]> : tensor<1xi64>
-+    %p5 = stablehlo.constant dense<[5]> : tensor<1xi64>
-+    %p6 = stablehlo.constant dense<[6]> : tensor<1xi64>
-+    %p7 = stablehlo.constant dense<[7]> : tensor<1xi64>
-+    %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) {
-+      programs=[[@all_gather], [@all_gather], [@all_gather], [@all_gather],
-+                [@all_gather], [@all_gather], [@all_gather], [@all_gather]]
-+    } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>,
-+         tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) ->
-+        (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>,
-+         tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>)
-+    check.expect_eq_const %results#0, dense<[0, 2]> : tensor<2xi64>
-+    check.expect_eq_const %results#1, dense<[1, 3]> : tensor<2xi64>
-+    check.expect_eq_const %results#2, dense<[0, 2]> : tensor<2xi64>
-+    check.expect_eq_const %results#3, dense<[1, 3]> : tensor<2xi64>
-+    check.expect_eq_const %results#4, dense<[4, 6]> : tensor<2xi64>
-+    check.expect_eq_const %results#5, dense<[5, 7]> : tensor<2xi64>
-+    check.expect_eq_const %results#6, dense<[4, 6]> : tensor<2xi64>
-+    check.expect_eq_const %results#7, dense<[5, 7]> : tensor<2xi64>
-+    func.return
-+  }
-+}
-diff --ruN a/stablehlo/stablehlo/tests/interpret/all_reduce.mlir b/stablehlo/stablehlo/tests/interpret/all_reduce.mlir
---- stablehlo/stablehlo/tests/interpret/all_reduce.mlir
-+++ stablehlo/stablehlo/tests/interpret/all_reduce.mlir
-@@ -135,3 +135,45 @@
-     func.return
-   }
- }
-+
-+// -----
-+
-+module @mesh_axes_subaxis {
-+  func.func @all_reduce(%operand : tensor<1xi64>) -> tensor<1xi64> {
-+    %result = "stablehlo.all_reduce"(%operand) ({
-+      ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
-+        %0 = stablehlo.add %arg0, %arg1 : tensor<i64>
-+        stablehlo.return %0 : tensor<i64>
-+    }) {
-+      replica_groups = #stablehlo.replica_group_mesh_axes<mesh = #stablehlo.mesh<axes = [#stablehlo.mesh_axis<name = "x", size = 2>, #stablehlo.mesh_axis<name = "y", size = 4>]>, axes = [#stablehlo.axis_ref<name = "y", sub_axis_info = (1)2>]>,
-+      channel_handle = #stablehlo.channel_handle<handle = 0, type = 0>
-+    } : (tensor<1xi64>) -> tensor<1xi64>
-+    return %result : tensor<1xi64>
-+  }
-+  func.func @main() {
-+    %p0 = stablehlo.constant dense<[10]> : tensor<1xi64>
-+    %p1 = stablehlo.constant dense<[20]> : tensor<1xi64>
-+    %p2 = stablehlo.constant dense<[30]> : tensor<1xi64>
-+    %p3 = stablehlo.constant dense<[40]> : tensor<1xi64>
-+    %p4 = stablehlo.constant dense<[50]> : tensor<1xi64>
-+    %p5 = stablehlo.constant dense<[60]> : tensor<1xi64>
-+    %p6 = stablehlo.constant dense<[70]> : tensor<1xi64>
-+    %p7 = stablehlo.constant dense<[80]> : tensor<1xi64>
-+    %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) {
-+      programs=[[@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce],
-+                [@all_reduce], [@all_reduce], [@all_reduce], [@all_reduce]]
-+    } : (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>,
-+         tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>) ->
-+        (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>,
-+         tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>)
-+    check.expect_eq_const %results#0, dense<[40]> : tensor<1xi64>
-+    check.expect_eq_const %results#1, dense<[60]> : tensor<1xi64>
-+    check.expect_eq_const %results#2, dense<[40]> : tensor<1xi64>
-+    check.expect_eq_const %results#3, dense<[60]> : tensor<1xi64>
-+    check.expect_eq_const %results#4, dense<[120]> : tensor<1xi64>
-+    check.expect_eq_const %results#5, dense<[140]> : tensor<1xi64>
-+    check.expect_eq_const %results#6, dense<[120]> : tensor<1xi64>
-+    check.expect_eq_const %results#7, dense<[140]> : tensor<1xi64>
-+    func.return
-+  }
-+}
-diff --ruN a/stablehlo/stablehlo/tests/interpret/all_to_all.mlir b/stablehlo/stablehlo/tests/interpret/all_to_all.mlir
---- stablehlo/stablehlo/tests/interpret/all_to_all.mlir
-+++ stablehlo/stablehlo/tests/interpret/all_to_all.mlir
-@@ -172,3 +172,43 @@
-     func.return %results#0, %results#1, %results#2, %results#3 : tensor<4x2xi64>, tensor<6x2xi32>, tensor<4x2xi64>, tensor<6x2xi32>
-   }
- }
-+
-+// -----
-+
-+module @mesh_axes_subaxis {
-+  func.func @all_to_all(%operand : tensor<2x1xi64>) -> tensor<1x2xi64> {
-+    %result = "stablehlo.all_to_all"(%operand) {
-+      split_dimension = 0 : i64,
-+      concat_dimension = 1 : i64,
-+      split_count = 2 : i64,
-+      replica_groups = #stablehlo.replica_group_mesh_axes<mesh = #stablehlo.mesh<axes = [#stablehlo.mesh_axis<name = "x", size = 2>, #stablehlo.mesh_axis<name = "y", size = 4>]>, axes = [#stablehlo.axis_ref<name = "y", sub_axis_info = (1)2>]>
-+    } : (tensor<2x1xi64>) -> tensor<1x2xi64>
-+    return %result : tensor<1x2xi64>
-+  }
-+  func.func @main() {
-+    %p0 = stablehlo.constant dense<[[1], [2]]> : tensor<2x1xi64>
-+    %p1 = stablehlo.constant dense<[[3], [4]]> : tensor<2x1xi64>
-+    %p2 = stablehlo.constant dense<[[10], [20]]> : tensor<2x1xi64>
-+    %p3 = stablehlo.constant dense<[[30], [40]]> : tensor<2x1xi64>
-+    %p4 = stablehlo.constant dense<[[5], [6]]> : tensor<2x1xi64>
-+    %p5 = stablehlo.constant dense<[[7], [8]]> : tensor<2x1xi64>
-+    %p6 = stablehlo.constant dense<[[50], [60]]> : tensor<2x1xi64>
-+    %p7 = stablehlo.constant dense<[[70], [80]]> : tensor<2x1xi64>
-+    %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) {
-+      programs=[[@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all],
-+                [@all_to_all], [@all_to_all], [@all_to_all], [@all_to_all]]
-+    } : (tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>,
-+         tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>, tensor<2x1xi64>) ->
-+        (tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>,
-+         tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>, tensor<1x2xi64>)
-+    check.expect_eq_const %results#0, dense<[[1, 10]]> : tensor<1x2xi64>
-+    check.expect_eq_const %results#1, dense<[[3, 30]]> : tensor<1x2xi64>
-+    check.expect_eq_const %results#2, dense<[[2, 20]]> : tensor<1x2xi64>
-+    check.expect_eq_const %results#3, dense<[[4, 40]]> : tensor<1x2xi64>
-+    check.expect_eq_const %results#4, dense<[[5, 50]]> : tensor<1x2xi64>
-+    check.expect_eq_const %results#5, dense<[[7, 70]]> : tensor<1x2xi64>
-+    check.expect_eq_const %results#6, dense<[[6, 60]]> : tensor<1x2xi64>
-+    check.expect_eq_const %results#7, dense<[[8, 80]]> : tensor<1x2xi64>
-+    func.return
-+  }
-+}
-diff --ruN a/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir b/stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir
---- stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir
-+++ stablehlo/stablehlo/tests/interpret/reduce_scatter.mlir
-@@ -90,3 +90,45 @@
-     func.return
-   }
- }
-+
-+// -----
-+
-+module @mesh_axes_subaxis {
-+  func.func @reduce_scatter(%operand : tensor<2xi64>) -> tensor<1xi64> {
-+    %result = "stablehlo.reduce_scatter"(%operand) ({
-+      ^bb0(%arg0: tensor<i64>, %arg1: tensor<i64>):
-+        %0 = stablehlo.add %arg0, %arg1 : tensor<i64>
-+        stablehlo.return %0 : tensor<i64>
-+    }) {
-+      scatter_dimension = 0 : i64,
-+      replica_groups = #stablehlo.replica_group_mesh_axes<mesh = #stablehlo.mesh<axes = [#stablehlo.mesh_axis<name = "x", size = 2>, #stablehlo.mesh_axis<name = "y", size = 4>]>, axes = [#stablehlo.axis_ref<name = "y", sub_axis_info = (1)2>]>
-+    } : (tensor<2xi64>) -> tensor<1xi64>
-+    return %result : tensor<1xi64>
-+  }
-+  func.func @main() {
-+    %p0 = stablehlo.constant dense<[1, 2]> : tensor<2xi64>
-+    %p1 = stablehlo.constant dense<[3, 4]> : tensor<2xi64>
-+    %p2 = stablehlo.constant dense<[10, 20]> : tensor<2xi64>
-+    %p3 = stablehlo.constant dense<[30, 40]> : tensor<2xi64>
-+    %p4 = stablehlo.constant dense<[5, 6]> : tensor<2xi64>
-+    %p5 = stablehlo.constant dense<[7, 8]> : tensor<2xi64>
-+    %p6 = stablehlo.constant dense<[50, 60]> : tensor<2xi64>
-+    %p7 = stablehlo.constant dense<[70, 80]> : tensor<2xi64>
-+    %results:8 = "interpreter.run_parallel"(%p0, %p1, %p2, %p3, %p4, %p5, %p6, %p7) {
-+      programs=[[@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter],
-+                [@reduce_scatter], [@reduce_scatter], [@reduce_scatter], [@reduce_scatter]]
-+    } : (tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>,
-+         tensor<2xi64>, tensor<2xi64>, tensor<2xi64>, tensor<2xi64>) ->
-+        (tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>,
-+         tensor<1xi64>, tensor<1xi64>, tensor<1xi64>, tensor<1xi64>)
-+    check.expect_eq_const %results#0, dense<[11]> : tensor<1xi64>
-+    check.expect_eq_const %results#1, dense<[33]> : tensor<1xi64>
-+    check.expect_eq_const %results#2, dense<[22]> : tensor<1xi64>
-+    check.expect_eq_const %results#3, dense<[44]> : tensor<1xi64>
-+    check.expect_eq_const %results#4, dense<[55]> : tensor<1xi64>
-+    check.expect_eq_const %results#5, dense<[77]> : tensor<1xi64>
-+    check.expect_eq_const %results#6, dense<[66]> : tensor<1xi64>
-+    check.expect_eq_const %results#7, dense<[88]> : tensor<1xi64>
-+    func.return
-+  }
-+}
 diff --ruN a/stablehlo/stablehlo/tests/ops_broadcasting.mlir b/stablehlo/stablehlo/tests/ops_broadcasting.mlir
 --- stablehlo/stablehlo/tests/ops_broadcasting.mlir
 +++ stablehlo/stablehlo/tests/ops_broadcasting.mlir
@@ -984,98 +505,6 @@ diff --ruN a/stablehlo/stablehlo/tests/ops_broadcasting.mlir b/stablehlo/stableh
 +  return %0 : tensor<3x4x5xf64>
 +}
 +
-diff --ruN a/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir b/stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir
---- stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir
-+++ stablehlo/stablehlo/tests/transforms/mesh_axes_replica_group_compatibility.mlir
-@@ -5,7 +5,7 @@
- 
-   // CHECK-LABEL: @all_reduce_rgv3
-   func.func @all_reduce_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> {
--    // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64>
-+    // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64>
-     %0 = "stablehlo.all_reduce"(%arg0) ({
-     ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-       %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-@@ -19,7 +19,7 @@
- 
-   // CHECK-LABEL: @all_gather_rgv3
-   func.func @all_gather_rgv3(%arg0: tensor<4xf32>) -> tensor<8xf32> {
--    // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64>
-+    // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64>
-     %0 = "stablehlo.all_gather"(%arg0) {
-       all_gather_dim = 0 : i64,
-       replica_groups = #stablehlo.replica_group_mesh_axes<mesh = #stablehlo.mesh<axes = [#stablehlo.mesh_axis<name = "x", size = 2>, #stablehlo.mesh_axis<name = "y", size = 2>]>, axes = [#stablehlo.axis_ref<name = "y">]>
-@@ -30,7 +30,7 @@
- 
-   // CHECK-LABEL: @all_to_all_rgv3
-   func.func @all_to_all_rgv3(%arg0: tensor<4xf32>) -> tensor<4xf32> {
--    // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64>
-+    // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64>
-     %0 = "stablehlo.all_to_all"(%arg0) {
-       concat_dimension = 0 : i64,
-       split_dimension = 0 : i64,
-@@ -44,7 +44,7 @@
- 
-   // CHECK-LABEL: @all_reduce_sdy_mesh
-   func.func @all_reduce_sdy_mesh(%arg0: tensor<4xf32>) -> tensor<4xf32> {
--    // CHECK: replica_groups = dense<{{\[\[}}0, 2], [1, 3]]> : tensor<2x2xi64>
-+    // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3]]> : tensor<2x2xi64>
-     %0 = "stablehlo.all_reduce"(%arg0) ({
-     ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-       %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-@@ -59,7 +59,7 @@
- 
-   // CHECK-LABEL: @all_reduce_sdy_mesh_dev
-   func.func @all_reduce_sdy_mesh_dev(%arg0: tensor<4xf32>) -> tensor<4xf32> {
--    // CHECK: replica_groups = dense<{{\[\[}}0, 1], [2, 3]]> : tensor<2x2xi64>
-+    // CHECK{LITERAL}: replica_groups = dense<[[0, 1], [2, 3]]> : tensor<2x2xi64>
-     %0 = "stablehlo.all_reduce"(%arg0) ({
-     ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-       %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-@@ -84,4 +84,43 @@
-     } : (tensor<4xf32>) -> tensor<4xf32>
-     return %0 : tensor<4xf32>
-   }
-+
-+  // CHECK-LABEL: @all_reduce_subaxis
-+  func.func @all_reduce_subaxis(%arg0: tensor<4xf32>) -> tensor<4xf32> {
-+    // CHECK{LITERAL}: replica_groups = dense<[[0, 2], [1, 3], [4, 6], [5, 7]]> : tensor<4x2xi64>
-+    %0 = "stablehlo.all_reduce"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+    }) {
-+      replica_groups = #stablehlo.replica_group_mesh_axes<mesh = #stablehlo.mesh<axes = [#stablehlo.mesh_axis<name = "x", size = 2>, #stablehlo.mesh_axis<name = "y", size = 4>]>, axes = [#stablehlo.axis_ref<name = "y", sub_axis_info = (1)2>]>
-+    } : (tensor<4xf32>) -> tensor<4xf32>
-+    return %0 : tensor<4xf32>
-+  }
-+
-+  // CHECK-LABEL: @all_reduce_subaxis_order_1
-+  func.func @all_reduce_subaxis_order_1(%arg0: tensor<4xf32>) -> tensor<4xf32> {
-+    // CHECK{LITERAL}: replica_groups = dense<[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [10, 11, 12, 13, 14, 15, 16, 17, 18, 19], [20, 21, 22, 23, 24, 25, 26, 27, 28, 29]]> : tensor<3x10xi64>
-+    %0 = "stablehlo.all_reduce"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+    }) {
-+      replica_groups = #stablehlo.replica_group_mesh_axes<mesh = #stablehlo.mesh<axes = [#stablehlo.mesh_axis<name = "foo", size = 6>, #stablehlo.mesh_axis<name = "bar", size = 5>]>, axes = [#stablehlo.axis_ref<name = "foo", sub_axis_info = (3)2>, #stablehlo.axis_ref<name = "bar">]>
-+    } : (tensor<4xf32>) -> tensor<4xf32>
-+    return %0 : tensor<4xf32>
-+  }
-+
-+  // CHECK-LABEL: @all_reduce_subaxis_order_2
-+  func.func @all_reduce_subaxis_order_2(%arg0: tensor<4xf32>) -> tensor<4xf32> {
-+    // CHECK{LITERAL}: replica_groups = dense<[[0, 5, 1, 6, 2, 7, 3, 8, 4, 9], [10, 15, 11, 16, 12, 17, 13, 18, 14, 19], [20, 25, 21, 26, 22, 27, 23, 28, 24, 29]]> : tensor<3x10xi64>
-+    %0 = "stablehlo.all_reduce"(%arg0) ({
-+    ^bb0(%arg1: tensor<f32>, %arg2: tensor<f32>):
-+      %1 = "stablehlo.add"(%arg1, %arg2) : (tensor<f32>, tensor<f32>) -> tensor<f32>
-+      "stablehlo.return"(%1) : (tensor<f32>) -> ()
-+    }) {
-+      replica_groups = #stablehlo.replica_group_mesh_axes<mesh = #stablehlo.mesh<axes = [#stablehlo.mesh_axis<name = "foo", size = 6>, #stablehlo.mesh_axis<name = "bar", size = 5>]>, axes = [#stablehlo.axis_ref<name = "bar">, #stablehlo.axis_ref<name = "foo", sub_axis_info = (3)2>]>
-+    } : (tensor<4xf32>) -> tensor<4xf32>
-+    return %0 : tensor<4xf32>
-+  }
- }
 diff --ruN a/stablehlo/stablehlo/tests/verify_convolution.mlir b/stablehlo/stablehlo/tests/verify_convolution.mlir
 --- stablehlo/stablehlo/tests/verify_convolution.mlir
 +++ stablehlo/stablehlo/tests/verify_convolution.mlir
diff --git a/third_party/xla/workspace1.bzl b/third_party/xla/workspace1.bzl
index 54b16631be4f72..ffcaf9ddb768c7 100644
--- a/third_party/xla/workspace1.bzl
+++ b/third_party/xla/workspace1.bzl
@@ -27,7 +27,8 @@ def workspace():
     llvm_setup(name = "llvm-project")
     native.register_toolchains("@local_config_python//:py_toolchain")
     rules_pkg_dependencies()
-    compatibility_proxy_repo()
+    if "cc_compatibility_proxy" not in native.existing_rules():
+        compatibility_proxy_repo()
 
     tf_http_archive(
         name = "bazel_toolchains",
diff --git a/third_party/xla/workspace3.bzl b/third_party/xla/workspace3.bzl
index 79a81c7a644602..a52a078a5f3840 100644
--- a/third_party/xla/workspace3.bzl
+++ b/third_party/xla/workspace3.bzl
@@ -32,9 +32,9 @@ def workspace():
     # https://github.com/bazelbuild/bazel-skylib/releases
     tf_http_archive(
         name = "bazel_skylib",
-        sha256 = "bc283cdfcd526a52c3201279cda4bc298652efa898b10b4db0837dc51652756f",
+        sha256 = "3b5b49006181f5f8ff626ef8ddceaa95e9bb8ad294f7b5d7b11ea9f7ddaf8c59",
         urls = tf_mirror_urls(
-            "https://github.com/bazelbuild/bazel-skylib/releases/download/1.7.1/bazel-skylib-1.7.1.tar.gz",
+            "https://github.com/bazelbuild/bazel-skylib/releases/download/1.9.0/bazel-skylib-1.9.0.tar.gz",
         ),
     )
 
@@ -63,12 +63,11 @@ def workspace():
 
     tf_http_archive(
         name = "rules_cc",
-        urls = tf_mirror_urls("https://github.com/bazelbuild/rules_cc/releases/download/0.2.0/rules_cc-0.2.0.tar.gz"),
-        strip_prefix = "rules_cc-0.2.0",
-        sha256 = "ae244f400218f4a12ee81658ff246c0be5cb02c5ca2de5519ed505a6795431e9",
-        patch_file = [
-            "@xla//third_party/py:rules_cc_protobuf.patch",
-        ],
+        sha256 = "69e05df29f0010ba248ef8dafc1f084c8fd2f5c553da634422d8167f5c4b277b",
+        strip_prefix = "rules_cc-0.2.20",
+        urls = tf_mirror_urls(
+            "https://github.com/bazelbuild/rules_cc/releases/download/0.2.20/rules_cc-0.2.20.tar.gz",
+        ),
     )
 
     # Toolchains for ML projects hermetic builds.
diff --git a/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc b/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc
index 1f1d6ff1a06ab3..4c0577acedec25 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/cudnn.cc
@@ -153,6 +153,11 @@ bool IsSupportedCudnnFusion(const HloInstruction& instr,
     return true;
   }
 
+  if (hero->shape().element_type() == PrimitiveType::F64) {
+    VLOG(1) << "cuDNN GEMM fusion does not support F64.";
+    return false;
+  }
+
   stream_executor::CudaComputeCapability compute_capability =
       target_config.device_description.cuda_compute_capability();
   if ((compute_capability.IsAtLeastAmpere() &&
diff --git a/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc b/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc
index ecba9d7af5786c..f6ab20f5dc74b5 100644
--- a/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc
+++ b/third_party/xla/xla/backends/gpu/autotuner/cudnn_test.cc
@@ -47,7 +47,9 @@ namespace gpu {
 
 using CudnnBackendConfig = stream_executor::dnn::AlgorithmProto;
 
+using ::absl_testing::IsOkAndHolds;
 using ::testing::Gt;
+using ::testing::IsEmpty;
 using ::testing::SizeIs;
 using ::tsl::proto_testing::EqualsProto;
 
@@ -104,6 +106,22 @@ absl::string_view kTritonGemmFusionHlo = R"hlo(
       backend_config={"fusion_backend_config": {kind: "__triton_gemm"}}
   })hlo";
 
+absl::string_view kF64GemmFusionHlo = R"hlo(
+  fusion1 {
+    p0 = f64[3,28,32] parameter(0)
+    p1 = f64[3,28,32] parameter(1)
+    ROOT d = f64[3,32,32] dot(p0, p1),
+      lhs_batch_dims={0}, rhs_batch_dims={0},
+      lhs_contracting_dims={1}, rhs_contracting_dims={1}
+  }
+
+  e {
+    p0 = f64[3,28,32] parameter(0)
+    p1 = f64[3,28,32] parameter(1)
+    ROOT _ = f64[3,32,32] fusion(p0, p1), kind=kCustom, calls=fusion1,
+      backend_config={"fusion_backend_config": {kind: "__triton_gemm"}}
+  })hlo";
+
 absl::string_view kScaledDotGemmFusionHlo = R"hlo(
   block_scaled_dot {
     lhs = f8e4m3fn[256,128] parameter(0)
@@ -215,6 +233,15 @@ TEST_F(CudnnBackendTest, GetSupportedConfigsFromTritonGemmFusion) {
   EXPECT_THAT(configs, absl_testing::IsOkAndHolds(SizeIs(Gt(0))));
 }
 
+TEST_F(CudnnBackendTest, GetSupportedConfigsFromF64GemmFusionReturnsEmpty) {
+  ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> hlo_module,
+                       ParseAndReturnVerifiedModule(kF64GemmFusionHlo));
+  absl::StatusOr<std::vector<std::unique_ptr<BackendConfig>>> configs =
+      backend_->GetSupportedConfigs(
+          (*hlo_module->entry_computation()->root_instruction()));
+  EXPECT_THAT(configs, IsOkAndHolds(IsEmpty()));
+}
+
 TEST_F(CudnnBackendTest, GetSupportedConfigsFromScaledDotGemmFusion) {
   se::CudaComputeCapability cc =
       stream_executor_->GetDeviceDescription().cuda_compute_capability();
diff --git a/third_party/xla/xla/backends/gpu/codegen/BUILD b/third_party/xla/xla/backends/gpu/codegen/BUILD
index 94c732c4dc10e1..0877e5e0921bf2 100644
--- a/third_party/xla/xla/backends/gpu/codegen/BUILD
+++ b/third_party/xla/xla/backends/gpu/codegen/BUILD
@@ -94,16 +94,13 @@ xla_test(
         "//xla/service:pattern_matcher",
         "//xla/service/gpu:cudnn_support_utils",
         "//xla/service/gpu:ir_emission_utils",
-        "//xla/service/gpu:stream_executor_util",
         "//xla/stream_executor:device_description",
-        "//xla/stream_executor:dnn",
         "//xla/stream_executor:platform_manager",
+        "//xla/stream_executor:semantic_version",
         "//xla/stream_executor:stream_executor_h",
         "//xla/stream_executor/cuda:cuda_compute_capability",
-        "//xla/tests:hlo_pjrt_interpreter_reference_mixin",
+        "//xla/tests:hlo_interpreter_reference_mixin",
         "//xla/tsl/platform:env",
-        "//xla/tsl/platform:errors",
-        "//xla/tsl/platform:statusor",
         "//xla/tsl/platform:test",
         "@com_google_absl//absl/log:check",
         "@com_google_absl//absl/status:status_macros",
diff --git a/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc b/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc
index 216528dd11e2f9..8f47dd1e52192c 100644
--- a/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc
+++ b/third_party/xla/xla/backends/gpu/codegen/cudnn_test.cc
@@ -44,18 +44,15 @@ limitations under the License.
 #include "xla/service/dump.h"
 #include "xla/service/gpu/cudnn_support_utils.h"
 #include "xla/service/gpu/ir_emission_utils.h"
-#include "xla/service/gpu/stream_executor_util.h"
 #include "xla/service/hlo_module_config.h"
 #include "xla/service/pattern_matcher.h"
 #include "xla/stream_executor/cuda/cuda_compute_capability.h"
 #include "xla/stream_executor/device_description.h"
-#include "xla/stream_executor/dnn.h"
 #include "xla/stream_executor/platform_manager.h"
+#include "xla/stream_executor/semantic_version.h"
 #include "xla/stream_executor/stream_executor.h"
-#include "xla/tests/hlo_pjrt_interpreter_reference_mixin.h"
+#include "xla/tests/hlo_interpreter_reference_mixin.h"
 #include "xla/tsl/platform/env.h"
-#include "xla/tsl/platform/errors.h"
-#include "xla/tsl/platform/statusor.h"
 #include "xla/tsl/platform/test.h"
 #include "xla/xla.pb.h"
 #include "xla/xla_data.pb.h"
@@ -1407,13 +1404,20 @@ TEST_F(CuDnnFusionRewriteTest,
   // With other backends disabled, compilation must fail.
   ASSERT_OK_AND_ASSIGN(std::unique_ptr<VerifiedHloModule> module,
                        ParseAndReturnVerifiedModule(R"(
-e {
+triton_gemm_dot {
   p0 = f64[20,40,64] parameter(0)
   p0n = f64[20,40,64] negate(p0)
   p1 = f64[20,80,64] parameter(1)
-  r = f64[20,40,80] dot(p0n, p1),
+  ROOT r = f64[20,40,80] dot(p0n, p1),
     lhs_batch_dims={0}, rhs_batch_dims={0},
     lhs_contracting_dims={2}, rhs_contracting_dims={2}
+}
+
+e {
+  p0 = f64[20,40,64] parameter(0)
+  p1 = f64[20,80,64] parameter(1)
+  ROOT fusion = f64[20,40,80] fusion(p0, p1), kind=kCustom, calls=triton_gemm_dot,
+    backend_config={"fusion_backend_config": {kind: "__triton_gemm"}}
 })"));
   auto status =
       CreateExecutable(std::move(module), /*run_hlo_passes=*/true).status();
diff --git a/tools/tf_env_collect.sh b/tools/tf_env_collect.sh
index 390b4b6fd18f1e..7d09b97887f3fc 100755
--- a/tools/tf_env_collect.sh
+++ b/tools/tf_env_collect.sh
@@ -18,8 +18,10 @@ set -u  # Check for undefined variables
 
 # Track temporary files so they are removed on exit, including on interrupt.
 LOADED_LIBS_FILE=""
+ACCEL_FLAGS_FILE=""
 cleanup() {
   [ -n "${LOADED_LIBS_FILE:-}" ] && rm -f "$LOADED_LIBS_FILE"
+  [ -n "${ACCEL_FLAGS_FILE:-}" ] && rm -f "$ACCEL_FLAGS_FILE"
 }
 trap cleanup EXIT INT TERM
 
@@ -79,6 +81,14 @@ case "${OUTPUT_FILE##*/}" in
   *)   JSON_FILE="${OUTPUT_FILE}.json" ;;
 esac
 
+# Only pay for a temp file when the JSON summary is actually requested; it's
+# used to smuggle a couple of accelerator-detection booleans out of the
+# report-generation subshell below so we don't have to re-probe nvidia-smi /
+# rocm-smi / tensorflow-metal a second time just for the JSON section.
+if [ "$EMIT_JSON" -eq 1 ]; then
+  ACCEL_FLAGS_FILE="$(mktemp 2>/dev/null || mktemp -t tfenv)"
+fi
+
 echo "Collecting system information..."
 
 PYTHON_BIN_PATH="$(command -v python || command -v python3 || die "Cannot find Python binary")"
@@ -93,11 +103,15 @@ have_cmd() {
 
 run_cmd() {
   # Run a command if it exists, otherwise note that it is missing instead of
-  # erroring out. Captures stderr so the report stays readable.
+  # erroring out. Captures stderr so the report stays readable. Propagates
+  # the real exit status of the command (or 127 if it wasn't found) so
+  # callers can check success without re-invoking the command.
   if have_cmd "$1"; then
     "$@" 2>&1
+    return $?
   else
     echo "$1 not found"
+    return 127
   fi
 }
 
@@ -125,8 +139,10 @@ pip_run() {
 TF_PKG_PATTERN='^(tensorflow|tf-nightly|tensorflow-cpu|tensorflow-gpu|tensorflow-rocm|tensorflow-macos|tensorflow-metal|intel-tensorflow)\b'
 
 HEADER_WIDTH=68
-# Create a string of HEADER_WIDTH "=" characters
-HEADER=$(printf "%*s" "$HEADER_WIDTH" "" | sed 's/ /=/g')
+# Build a string of HEADER_WIDTH "=" characters using shell builtins only
+# (printf -v + parameter expansion), avoiding a fork+pipe through sed.
+printf -v HEADER '%*s' "$HEADER_WIDTH" ''
+HEADER=${HEADER// /=}
 
 print_header () {
   # This function simply prints the header with even spacing,
@@ -212,8 +228,12 @@ EOF
     echo "Not found"
   fi
 
+  # Fetch "pip list" once and reuse it below for the TensorFlow package
+  # conflict check, instead of shelling out to pip a second time.
+  PIP_LIST_OUTPUT="$(pip_run list 2>&1)"
+
   print_header 'check pips'
-  pip_run list 2>&1 | grep -E 'proto|numpy|keras|tensorflow|tf_nightly|tf-nightly'
+  grep -E 'proto|numpy|keras|tensorflow|tf_nightly|tf-nightly' <<<"$PIP_LIST_OUTPUT"
 
 
   print_header 'check for virtualenv'
@@ -232,12 +252,12 @@ EOF
   print_header 'tensorflow package conflicts'
   # Multiple TensorFlow distributions in the same environment frequently cause
   # confusing import errors; surface them so triage can spot the conflict.
-  TF_PKGS="$(pip_run list 2>/dev/null | grep -iE "$TF_PKG_PATTERN" || true)"
+  TF_PKGS="$(grep -iE "$TF_PKG_PATTERN" <<<"$PIP_LIST_OUTPUT" || true)"
   if [ -z "$TF_PKGS" ]; then
     echo "No TensorFlow packages found via pip."
   else
     echo "$TF_PKGS"
-    TF_COUNT="$(echo "$TF_PKGS" | grep -icE "$TF_PKG_PATTERN")"
+    TF_COUNT="$(grep -icE "$TF_PKG_PATTERN" <<<"$TF_PKGS")"
     if [ "$TF_COUNT" -gt 1 ]; then
       echo "WARNING: multiple TensorFlow distributions detected; this can cause import conflicts."
     fi
@@ -315,16 +335,17 @@ EOF
 
   print_header 'build / hermetic accelerator config'
   # Surface the environment variables that control modern (hermetic) CUDA and
-  # ROCm builds. See .bazelrc for how these are consumed.
+  # ROCm builds. See .bazelrc for how these are consumed. Uses indirect
+  # parameter expansion instead of eval - one less string re-parse per
+  # variable, and no eval footgun.
   for var in CC CXX \
              TF_NEED_CUDA TF_NEED_ROCM TF_CUDA_VERSION TF_CUDNN_VERSION \
              HERMETIC_CUDA_VERSION HERMETIC_CUDNN_VERSION \
              CUDA_HOME CUDA_PATH CUDA_TOOLKIT_PATH \
              ROCM_PATH HIP_PATH \
              XLA_FLAGS TF_XLA_FLAGS TPU_NAME; do
-    eval "marker=\${$var+set} val=\"\$$var\""
-    if [ "${marker:-}" = "set" ]; then
-      echo "$var=$val"
+    if [ -n "${!var+set}" ]; then
+      echo "$var=${!var}"
     else
       echo "$var is unset"
     fi
@@ -332,6 +353,10 @@ EOF
 
   print_header 'accelerator: nvidia gpu'
   run_cmd nvidia-smi
+  NVIDIA_STATUS=$?
+  if [ -n "$ACCEL_FLAGS_FILE" ] && [ "$NVIDIA_STATUS" -eq 0 ]; then
+    echo "HAS_NVIDIA=1" >> "$ACCEL_FLAGS_FILE"
+  fi
 
   print_header 'cuda libs'
   # Find cudart/cudnn files
@@ -343,6 +368,10 @@ EOF
 
   print_header 'accelerator: amd / rocm gpu'
   run_cmd rocm-smi
+  ROCM_STATUS=$?
+  if [ -n "$ACCEL_FLAGS_FILE" ] && [ "$ROCM_STATUS" -eq 0 ]; then
+    echo "HAS_ROCM=1" >> "$ACCEL_FLAGS_FILE"
+  fi
   if [ "$VERBOSE" -eq 1 ]; then
     print_header 'rocminfo'
     run_cmd rocminfo
@@ -355,9 +384,16 @@ EOF
   # tensorflow-metal is the PluggableDevice that enables GPU acceleration on
   # Apple Silicon; report whether it is installed and the GPU chipset.
   if [ "$(uname -s)" = "Darwin" ]; then
-    if pip_run show tensorflow-metal >/dev/null 2>&1; then
+    # Single "pip show" call, reused both for the existence check and for
+    # the Name/Version detail line below (previously called twice).
+    TF_METAL_INFO="$(pip_run show tensorflow-metal 2>&1)"
+    TF_METAL_STATUS=$?
+    if [ "$TF_METAL_STATUS" -eq 0 ]; then
       echo "tensorflow-metal installed:"
-      pip_run show tensorflow-metal 2>&1 | grep -iE '^(Name|Version):'
+      grep -iE '^(Name|Version):' <<<"$TF_METAL_INFO"
+      if [ -n "$ACCEL_FLAGS_FILE" ]; then
+        echo "HAS_METAL=1" >> "$ACCEL_FLAGS_FILE"
+      fi
     else
       echo "tensorflow-metal not installed"
     fi
@@ -396,14 +432,23 @@ EOF
 # Optional machine-readable JSON summary
 # ----------------------------------------------------------------------------
 if [ "$EMIT_JSON" -eq 1 ]; then
-  # Detect accelerators in the shell and hand the booleans to Python, which
-  # assembles a structured, easy-to-parse summary of the key facts.
+  # nvidia-smi / rocm-smi / tensorflow-metal detection already happened once
+  # above (inside the report-generation block); read the results back in
+  # from ACCEL_FLAGS_FILE instead of re-invoking those (potentially slow)
+  # tools a second time.
   HAS_NVIDIA=0
-  if have_cmd nvidia-smi && nvidia-smi >/dev/null 2>&1; then HAS_NVIDIA=1; fi
   HAS_ROCM=0
-  if have_cmd rocm-smi && rocm-smi >/dev/null 2>&1; then HAS_ROCM=1; fi
   HAS_METAL=0
-  if pip_run show tensorflow-metal >/dev/null 2>&1; then HAS_METAL=1; fi
+  if [ -n "$ACCEL_FLAGS_FILE" ] && [ -s "$ACCEL_FLAGS_FILE" ]; then
+    # shellcheck disable=SC1090
+    . "$ACCEL_FLAGS_FILE"
+  fi
+  # The block above only probes tensorflow-metal on Darwin hosts (matching
+  # the text report). On any other platform, fall back to a direct check so
+  # the JSON summary's accelerator info stays complete.
+  if [ "$HAS_METAL" -ne 1 ] && [ "$(uname -s)" != "Darwin" ]; then
+    if pip_run show tensorflow-metal >/dev/null 2>&1; then HAS_METAL=1; fi
+  fi
 
   TFENV_HAS_NVIDIA="$HAS_NVIDIA" \
   TFENV_HAS_ROCM="$HAS_ROCM" \